Jacfwd

Contents

Jacfwd#

nabla.transforms.jacfwd.jacfwd(func, argnums=0, has_aux=False, holomorphic=False, allow_int=False)[source]#

Prototype implementation of jacfwd using forward-mode autodiff.

This computes the Jacobian using the pattern: vmap(jvp(func, primals, tangents), in_axes=(primal_axes, tangent_axes))

where primal_axes are None (broadcast) and tangent_axes are 0 (vectorize).

Parameters:

  • func (Callable[[...], Any]) – Function to differentiate

  • argnums (int | tuple[int, ...] | list[int]) – Which arguments to differentiate with respect to

  • has_aux (bool) – Whether function returns auxiliary data

  • holomorphic (bool) – Ignored (for JAX compatibility)

  • allow_int (bool) – Ignored (for JAX compatibility)

Returns:

Function that computes the Jacobian using forward-mode autodiff

Return type:

Callable[[…], Any]

Contents