Derivatives

Selective coordinate-derivative utilities avoid materializing full Jacobians and Hessians that a PDE does not use.

The second argument, argnums, contains positional indices of scalar arguments at the single-point function boundary:

def scalar_prediction(model_state, t, x):
    return model_apply(model_state, jnp.stack((t, x)))[0]


value_and_first = value_and_jacobian(scalar_prediction, (1, 2))
value, (du_dt, du_dx) = value_and_first(model_state, t, x)

second_spatial = hessian_diagonal(scalar_prediction, 2)
(d2u_dx2,) = second_spatial(model_state, t, x)

The derivative tuple preserves requested index order. Batch the transformed single-point function with jax.vmap.

value_and_jacobian [source]

value_and_jacobian(
    function: Callable[..., Any], argnums: ArgumentNumbers
) -> Callable[..., tuple[Any, tuple[Any, ...]]]

Build a forward-mode value-and-Jacobian function for selected scalar arguments.

The returned callable linearizes one primal evaluation and applies one JVP per selected argument. Each selected argument must be a scalar array. PINN equations can therefore split a point into scalar coordinates, call this helper at one point, and apply :func:jax.vmap outside it for collocation batches.

Parameters:

Name	Type	Description	Default
function	Callable[..., Any]	Pure function differentiated with respect to selected scalar positional arguments.	required
argnums	ArgumentNumbers	Unique positional argument number or numbers, ordered like the requested Jacobian entries.	required

Returns:

Type	Description
Callable[..., tuple[Any, tuple[Any, ...]]]	Callable returning the primal output and one derivative PyTree per selected argument.

Raises:

Type	Description
TypeError	If argnums contains non-integer values or Boolean values.
ValueError	If argnums is empty or duplicated, or a selected primal argument is not scalar.

hessian_diagonal [source]

hessian_diagonal(
    function: Callable[..., Any], argnums: ArgumentNumbers
) -> Callable[..., tuple[Any, ...]]

Build forward-mode pure second derivatives for selected scalar arguments.

Unlike a complete Hessian, this helper computes only d²function / darg² for each selected argument. It avoids tracing unused mixed derivatives in PDE operators such as one-dimensional diffusion and coordinate Laplacians. Every selected argument must be a scalar array.

Parameters:

Name	Type	Description	Default
function	Callable[..., Any]	Pure function differentiated with respect to selected scalar positional arguments.	required
argnums	ArgumentNumbers	Unique positional argument number or numbers, ordered like the requested Hessian diagonal entries.	required

Returns:

Type	Description
Callable[..., tuple[Any, ...]]	Callable returning one pure-second-derivative PyTree per selected argument.

Raises:

Type	Description
TypeError	If argnums contains non-integer values or Boolean values.
ValueError	If argnums is empty or duplicated, or a selected primal argument is not scalar.