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

DOC: update sharp bits info about fori_loop differentiability

Browse Source
This commit is contained in:
Jake VanderPlas 2022-02-25 09:15:55 -08:00
parent d5a1c64d13
commit 74e7fdfd4c
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/notebooks/Common_Gotchas_in_JAX.ipynb
View File

@ -1579,7 +1579,7 @@
"\n",
" - `lax.cond` _differentiable_\n",
" - `lax.while_loop` __fwd-mode-differentiable__\n",
" - `lax.fori_loop` __fwd-mode-differentiable__\n",
" - `lax.fori_loop` __fwd-mode-differentiable__ in general; __fwd and rev-mode differentiable__ if endpoints are static.\n",
" - `lax.scan` _differentiable_"
]
},

2
docs/notebooks/Common_Gotchas_in_JAX.md
View File

@ -811,7 +811,7 @@ There are more options for control flow in JAX. Say you want to avoid re-compila
- `lax.cond` _differentiable_
- `lax.while_loop` __fwd-mode-differentiable__
- `lax.fori_loop` __fwd-mode-differentiable__
- `lax.fori_loop` __fwd-mode-differentiable__ in general; __fwd and rev-mode differentiable__ if endpoints are static.
- `lax.scan` _differentiable_
+++ {"id": "Sd9xrLMXeK3A"}