From 74e7fdfd4c5e174d1dc0289400ea115752c98380 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Fri, 25 Feb 2022 09:15:55 -0800 Subject: [PATCH] DOC: update sharp bits info about fori_loop differentiability --- docs/notebooks/Common_Gotchas_in_JAX.ipynb | 2 +- docs/notebooks/Common_Gotchas_in_JAX.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/notebooks/Common_Gotchas_in_JAX.ipynb b/docs/notebooks/Common_Gotchas_in_JAX.ipynb index ce8c56debbba..ffc3540d4dce 100644 --- a/docs/notebooks/Common_Gotchas_in_JAX.ipynb +++ b/docs/notebooks/Common_Gotchas_in_JAX.ipynb @@ -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_" ] }, diff --git a/docs/notebooks/Common_Gotchas_in_JAX.md b/docs/notebooks/Common_Gotchas_in_JAX.md index d38838190384..b45939761428 100644 --- a/docs/notebooks/Common_Gotchas_in_JAX.md +++ b/docs/notebooks/Common_Gotchas_in_JAX.md @@ -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"}