[nnx] add custom_vjp to docs

google · Oct 8, 2024 · 312368e · 312368e
1 parent 40b9b3a
commit 312368e
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