adv_jax_math.sparse_pullback_map

adv_jax_math.sparse_pullback_map(fn, y, *, higher_order=False)[source]

Wrapper for sparsity exploiting pullback.

Wraps the given map with logic to ensure cotangents flow through the diagonal of its pullback. The derivatives will be exact for maps whose Jacobians are block diagonal.

See also

sparse_pullback

Applies the same transformation and immediatly returns its output.

Parameters:

  • fn (callable) – Vectorized map.

  • y (pytree) – Example input used to closure-convert fn.

  • higher_order (bool) – Whether to support higher-order differentiation with the HiJAX backend at the expense of evaluating the primal fn twice. The custom-VJP backend supports higher-order differentiation regardless of this flag. Default is False.

Returns:

wrapper (callable) – Same forward map but with a sparsity exploiting pullback.

Examples

>>> fn = sparse_pullback_map(fn, y)
>>> out = fn(y)