Improve documentation for jax.jacobian

jax-ml · Sep 26, 2024 · cf51ee7 · cf51ee7
1 parent 5d4cae0
commit cf51ee7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/docs/jax.rst b/docs/jax.rst
@@ -91,6 +91,7 @@ Automatic differentiation
 
     grad
     value_and_grad
+    jacobian
     jacfwd
     jacrev
     hessian

diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -680,7 +680,13 @@ def jacfun(*args, **kwargs):
       return jac_tree, aux
 
   return jacfun
-jacobian = jacrev
+
+
+def jacobian(fun: Callable, argnums: int | Sequence[int] = 0,
+             has_aux: bool = False, holomorphic: bool = False, allow_int: bool = False) -> Callable:
+  """Alias of :func:`jax.jacrev`."""
+  return jacrev(fun, argnums=argnums, has_aux=has_aux, holomorphic=holomorphic, allow_int=allow_int)
+
 
 _check_input_dtype_jacrev = partial(_check_input_dtype_revderiv, "jacrev")
 _check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev")