Jacobian

asdex.jacobian(f, input_shape, *, mode=None, symmetric=False)

Detect sparsity, color, and return a function computing sparse Jacobians.

Combines jacobian_coloring and jacobian_from_coloring in one call.

Parameters:

Name	Type	Description	Default
f	Callable[[ArrayLike], ArrayLike]	Function taking an array and returning an array. Input and output may be multi-dimensional.	required
input_shape	int | tuple[int, ...]	Shape of the input array.	required
mode	JacobianMode | None	AD mode. "fwd" uses JVPs (forward-mode AD), "rev" uses VJPs (reverse-mode AD). None picks whichever of fwd/rev needs fewer colors.	None
symmetric	bool	Whether to use symmetric (star) coloring. Requires a square Jacobian.	False

Returns:

Type	Description
Callable[[ArrayLike], BCOO]	A function that takes an input array and returns the sparse Jacobian as BCOO of shape (m, n) where n = x.size and m = prod(output_shape).

asdex.jacobian_from_coloring(f, coloring)

Build a sparse Jacobian function from a pre-computed coloring.

Uses row coloring + VJPs or column coloring + JVPs, depending on which needs fewer colors.

Parameters:

Name	Type	Description	Default
f	Callable[[ArrayLike], ArrayLike]	Function taking an array and returning an array. Input and output may be multi-dimensional.	required
coloring	ColoredPattern	Pre-computed ColoredPattern from jacobian_coloring.	required

Returns:

Type	Description
Callable[[ArrayLike], BCOO]	A function that takes an input array and returns the sparse Jacobian as BCOO of shape (m, n) where n = x.size and m = prod(output_shape).

asdex.jacobian_coloring(f, input_shape, *, mode=None, symmetric=False)

Detect Jacobian sparsity and color in one step.

Parameters:

Name	Type	Description	Default
f	Callable	Function taking an array and returning an array.	required
input_shape	int | tuple[int, ...]	Shape of the input array.	required
mode	JacobianMode | None	AD mode. "fwd" uses JVPs (forward-mode AD), "rev" uses VJPs (reverse-mode AD), None picks whichever of fwd/rev needs fewer colors (unless symmetric is True, in which case defaults to "fwd").	None
symmetric	bool	Whether to use symmetric (star) coloring. Requires a square Jacobian.	False

Returns:

Type	Description
ColoredPattern	A ColoredPattern ready for jacobian_from_coloring.

asdex.jacobian_coloring_from_sparsity(sparsity, *, mode=None, symmetric=False)

Color a sparsity pattern for sparse Jacobian computation.

Assigns colors so that same-colored rows (or columns) can be computed together in a single VJP (or JVP).

Parameters:

Name	Type	Description	Default
sparsity	SparsityPattern	Sparsity pattern of shape (m, n).	required
mode	JacobianMode | None	AD mode. "fwd" uses JVPs (column coloring), "rev" uses VJPs (row coloring). None picks whichever of fwd/rev needs fewer colors (unless symmetric is True, in which case defaults to "fwd").	None
symmetric	bool	Whether to use symmetric (star) coloring. Requires a square pattern.	False

Returns:

Type	Description
ColoredPattern	A ColoredPattern ready for jacobian_from_coloring.

asdex.jacobian_sparsity(f, input_shape)

Detect global Jacobian sparsity pattern for f: R^n -> R^m.

Analyzes the computation graph structure directly, without evaluating any derivatives. The result is valid for all inputs.

Parameters:

Name	Type	Description	Default
f	Callable	Function taking an array and returning an array.	required
input_shape	int | tuple[int, ...]	Shape of the input array. An integer is treated as a 1D length.	required

Returns:

Type	Description
SparsityPattern	SparsityPattern of shape (m, n) where n = prod(input_shape) and m = prod(output_shape). Entry (i, j) is present if output i depends on input j.

asdex.check_jacobian_correctness(f, x, coloring, *, method='matvec', num_probes=25, seed=0, rtol=None, atol=None)

Verify asdex's sparse Jacobian against a JAX reference at a given input.

Parameters:

Name	Type	Description	Default
f	Callable[[ArrayLike], ArrayLike]	Function taking an array and returning an array.	required
x	ArrayLike	Input at which to evaluate the Jacobian.	required
coloring	ColoredPattern	Pre-computed colored pattern from :func:~asdex.jacobian_coloring.	required
method	Literal['matvec', 'dense']	Verification method. "matvec" uses randomized matrix-vector products, which is O(k) in the number of probes. "dense" materializes the full dense Jacobian, which is O(n^2).	'matvec'
num_probes	int	Number of random probe vectors (only used by "matvec").	25
seed	int	PRNG seed for reproducibility (only used by "matvec").	0
rtol	float | None	Relative tolerance for comparison. Defaults to 1e-5 for "matvec" and 1e-7 for "dense".	None
atol	float | None	Absolute tolerance for comparison. Defaults to 1e-5 for "matvec" and 1e-7 for "dense".	None

Raises:

Type	Description
VerificationError	If the sparse and reference Jacobians disagree.

asdex.VerificationError

Bases: AssertionError

Raised when asdex's sparse result does not match JAX's dense reference.

This indicates that the detected sparsity pattern is missing nonzeros, which is a bug — asdex's patterns should always be conservative (i.e., contain at least all true nonzeros). If you encounter this error, please help out asdex's development by reporting this at https://github.com/adrhill/asdex/issues.

---

asdex.JacobianMode = Literal['fwd', 'rev'] module-attribute

AD mode for Jacobian computation.

"fwd" uses JVPs (forward-mode AD), "rev" uses VJPs (reverse-mode AD).