adv_jax_math.sparse_pullback

adv_jax_math.sparse_pullback(fn, y, /, batch_size=None, *, reduction=None, chunk_reduction=<function identity>, strip_dim0=False, shard=False, mesh=None, higher_order=False, **kwargs)[source]

Compute chunk_reduction(fn(fun_input)) in batches with sparse pullbacks.

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.

Notes

This method does not automatically wrap fn with vmap. Unless fn is already wrapped with vmap, the leading dimension of y will not be stripped before it is passed into fn. This can be inconvenient for nesting calls to sparse_pullback, since only batching along the first axis is currently supported. However, the strip_dim0 flag should cover the most common case of nesting calls where batch_size is one on the outermost call.

See also

sparse_pullback_map

Functional version.

Parameters:

  • fn (callable) – Vectorized map.

  • y (pytree) – Data to split into batches to feed to fn.

  • batch_size (int or None) – Size of batches. If no batching should be done or the batch size is the full input then supply None. chunk_size is accepted as an alias when batch_size is None.

  • reduction (callable or None) – Binary reduction operation. Should take two arguments and return one output, e.g. jnp.add.

  • chunk_reduction (callable) – Chunk-wise reduction operation. Should typically apply reduction along the mapped axis, e.g. jnp.add.reduce.

  • strip_dim0 (bool) – Whether to strip the leading dim of y before passing it to fn; see notes. This flag only works if batch_size is one. It should be set to False if fn is wrapped in vmap. Default is False.

  • shard (bool) – Whether to shard y across devices before applying chunked batching. The divisible prefix is split across devices; when supplied, batch_size bounds the batches processed on each device. A local remainder is evaluated once per device, and a final global remainder is evaluated once overall. The input length need not be divisible by either the device count or batch_size. Default is False.

  • mesh (jax.sharding.Mesh or None) – Optional one-dimensional mesh with an AxisType.Auto axis to use when shard=True. Default is None, which uses all available devices.

  • 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:

out (pytree) – Returns chunk_reduction(fn(y)).

Examples

>>> out = sparse_pullback(fn, y)