adv_jax_math.batch_jacrev

adv_jax_math.batch_jacrev(fun, argnums=0, has_aux=False, holomorphic=False, allow_int=False, *, batch_size=None, **kwargs)

Jacobian of fun evaluated row-by-row using reverse-mode AD.

Refrences

The original copyright notice is as follows Copyright 2018 The JAX Authors. Licensed under the Apache License, Version 2.0 (the “License”); https://github.com/jax-ml/jax/blob/main/jax/_src/api.py.

param fun:

Function whose Jacobian is to be computed.

type fun:

callable

param argnums:

Specifies which positional argument(s) to differentiate with respect to (default 0).

type argnums:

Optional, integer or sequence of integers.

param has_aux:

Indicates whether fun returns a pair where the first element is considered the output of the mathematical function to be differentiated and the second element is auxiliary data. Default False.

type has_aux:

Optional, bool.

param holomorphic:

Indicates whether fun is promised to be holomorphic. Default False.

type holomorphic:

Optional, bool.

param allow_int:

Whether to allow differentiating with respect to integer valued inputs. The gradient of an integer input will have a trivial vector-space dtype (float0). Default False.

type allow_int:

Optional, bool.

param batch_size:

The size of the batches to pass to vmap. If None, defaults to the largest possible batch_size. chunk_size is accepted as an alias when batch_size is None.

type batch_size:

int

returns:

jac (callable) – A function with the same arguments as fun, that evaluates the Jacobian of fun using reverse-mode automatic differentiation. If has_aux is True then a pair of (jacobian, auxiliary_data) is returned.