mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix up previous jaxpr.rst commit. (#2647)

Browse Source
This commit is contained in:
Skye Wanderman-Milne 2020-04-08 11:29:02 -07:00 committed by GitHub
parent 5174f6d5dc
commit f37f235183
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
docs/jaxpr.rst
View File

@ -1,4 +1,4 @@
Understanding jaxprs
Understanding Jaxprs
====================
Updated: February 14, 2020 (for commit 9e6fe64).
@ -373,18 +373,6 @@ For the example consider the function ``func11`` below::
return lax.scan(body, 0., (arr, ones))
print(api.make_jaxpr(func11)(onp.ones(16), 5.))
{ lambda c ; a b.
let d e = scan[ forward=True
jaxpr={ lambda ; a b c d e.
let f = mul c e
g = add b f
h = add g a
in (h, b) }
length=16
linear=(False, False, False, True, False)
num_carry=1
num_consts=1 ] b 0.0 a * c
in (d, e) }
{ lambda c ; a b.
let d e = scan[ forward=True
jaxpr={ lambda ; f a b c.