Merge pull request #4266 from google:custom-vjp-docs

PiperOrigin-RevId: 683619304
google · Oct 8, 2024 · 601a915 · 601a915
2 parents 62a3414 + 312368e
commit 601a915
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/docs_nnx/api_reference/flax.nnx/transforms.rst b/docs_nnx/api_reference/flax.nnx/transforms.rst
@@ -21,3 +21,4 @@ transforms
 .. autofunction:: vmap
 .. autofunction:: eval_shape
 .. autofunction:: cond
+.. autofunction:: custom_vjp