Full API¶

`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.hessian(f, input_shape, *, mode=None, symmetric=True)` ¶

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

Combines hessian_coloring and hessian_from_coloring in one call.

If f returns a squeezable shape like (1,) or (1, 1), it is automatically squeezed to scalar.

Parameters:

Name	Type	Description	Default
`f`	`Callable[[ArrayLike], ArrayLike]`	Scalar-valued function taking an array. Input may be multi-dimensional.	required
`input_shape`	`int \| tuple[int, ...]`	Shape of the input array.	required
`mode`	`HessianMode \| None`	AD composition strategy for Hessian-vector products. `"fwd_over_rev"` uses forward-over-reverse, `"rev_over_fwd"` uses reverse-over-forward, `"rev_over_rev"` uses reverse-over-reverse. Defaults to `"fwd_over_rev"`.	`None`
`symmetric`	`bool`	Whether to use symmetric (star) coloring. Defaults to True (exploits H = H^T for fewer colors).	`True`

Returns:

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

`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.hessian_from_coloring(f, coloring)` ¶

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

Uses symmetric (star) coloring and Hessian-vector products by default.

If f returns a squeezable shape like (1,) or (1, 1), it is automatically squeezed to scalar.

Parameters:

Name	Type	Description	Default
`f`	`Callable[[ArrayLike], ArrayLike]`	Scalar-valued function taking an array. Input may be multi-dimensional.	required
`coloring`	`ColoredPattern`	Pre-computed `ColoredPattern` from `hessian_coloring`.	required

Returns:

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

`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.hessian_coloring(f, input_shape, *, mode=None, symmetric=True)` ¶

Detect Hessian sparsity and color in one step.

Parameters:

Name	Type	Description	Default
`f`	`Callable`	Scalar-valued function taking an array.	required
`input_shape`	`int \| tuple[int, ...]`	Shape of the input array.	required
`mode`	`HessianMode \| None`	AD composition strategy for Hessian-vector products. `"fwd_over_rev"` uses forward-over-reverse, `"rev_over_fwd"` uses reverse-over-forward, `"rev_over_rev"` uses reverse-over-reverse. Defaults to `"fwd_over_rev"`.	`None`
`symmetric`	`bool`	Whether to use symmetric (star) coloring. Defaults to True (exploits H = H^T for fewer colors).	`True`

Returns:

Type	Description
`ColoredPattern`	A `ColoredPattern` ready for `hessian_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.hessian_coloring_from_sparsity(sparsity, *, mode=None, symmetric=True)` ¶

Color a sparsity pattern for sparse Hessian computation.

Parameters:

Name	Type	Description	Default
`sparsity`	`SparsityPattern`	Sparsity pattern of shape (n, n).	required
`mode`	`HessianMode \| None`	AD composition strategy for Hessian-vector products. `"fwd_over_rev"` uses forward-over-reverse, `"rev_over_fwd"` uses reverse-over-forward, `"rev_over_rev"` uses reverse-over-reverse. Defaults to `"fwd_over_rev"`.	`None`
`symmetric`	`bool`	Whether to use symmetric (star) coloring. Defaults to True (exploits Hessian symmetry for fewer colors).	`True`

Returns:

Type	Description
`ColoredPattern`	A `ColoredPattern` ready for `hessian_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.hessian_sparsity(f, input_shape)` ¶

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

Analyzes the Jacobian sparsity of the gradient function, without evaluating any derivatives. The result is valid for all inputs.

If f returns a squeezable shape like (1,) or (1, 1), it is automatically squeezed to scalar.

Parameters:

Name	Type	Description	Default
`f`	`Callable`	Scalar-valued function taking 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 `(n, n)` where `n = prod(input_shape)`. Entry `(i, j)` is present if `H[i, j]` may be nonzero.

`asdex.SparsityPattern` `dataclass` ¶

Sparse matrix pattern storing only structural information (no values).

Stores row and column indices separately for efficient access by the coloring and decompression stages.

Attributes:

Name	Type	Description
`rows`	`NDArray[int32]`	Row indices of non-zero entries, shape `(nnz,)`
`cols`	`NDArray[int32]`	Column indices of non-zero entries, shape `(nnz,)`
`shape`	`tuple[int, int]`	Matrix dimensions `(m, n)`
`input_shape`	`tuple[int, ...] \| None`	Shape of the function input that produced this pattern. Defaults to `(n,)` if not specified.

`nnz` `property` ¶

Number of non-zero elements.

`m` `property` ¶

Number of rows.

`n` `property` ¶

Number of columns.

`density` `property` ¶

Fraction of non-zero entries.

`col_to_rows` `cached` `property` ¶

Mapping from column index to list of row indices with non-zeros.

Used by the coloring algorithm to build the row conflict graph.

`row_to_cols` `cached` `property` ¶

Mapping from row index to list of column indices with non-zeros.

Used by the coloring algorithm to build the column conflict graph.

`__post_init__()` ¶

Validate inputs and set defaults.

`from_coo(rows, cols, shape, *, input_shape=None)` `classmethod` ¶

Create pattern from row and column index arrays.

Parameters:

Name	Type	Description	Default
`rows`	`NDArray[int32] \| list[int]`	Row indices of non-zero entries.	required
`cols`	`NDArray[int32] \| list[int]`	Column indices of non-zero entries.	required
`shape`	`tuple[int, int]`	Matrix dimensions `(m, n)`.	required
`input_shape`	`tuple[int, ...] \| None`	Shape of the function input. Defaults to `(n,)` if not specified.	`None`

`from_bcoo(bcoo)` `classmethod` ¶

Create pattern from JAX BCOO sparse matrix.

`from_dense(dense)` `classmethod` ¶

Create pattern from dense boolean/numeric matrix.

Non-zero entries indicate pattern positions.

`to_bcoo(data=None)` ¶

Convert to JAX BCOO sparse matrix.

Parameters:

Name	Type	Description	Default
`data`	`ndarray \| None`	Optional data values. If None, uses all 1s.	`None`

`todense()` ¶

Convert to dense numpy array with 1s at pattern positions.

`save(path)` ¶

Save sparsity pattern to an .npz file.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination file path.	required

`load(path)` `classmethod` ¶

Load sparsity pattern from an .npz file.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source file path.	required

`str()` ¶

Render sparsity pattern with header and dot/braille grid.

`repr()` ¶

Return compact single-line representation.

`asdex.ColoredPattern` `dataclass` ¶

Result of a graph coloring for sparse differentiation.

Attributes:

Name	Type	Description
`sparsity`	`SparsityPattern`	The sparsity pattern that was colored.
`colors`	`NDArray[int32]`	Color assignment array. Shape `(m,)` for `"rev"` mode, `(n,)` for all other modes.
`num_colors`	`int`	Total number of colors used.
`symmetric`	`bool`	Whether symmetric (star) coloring was used.
`mode`	`ColoringMode`	The AD mode. Resolved, never `"auto"`. `"fwd"` uses JVPs (forward-mode AD), `"rev"` uses VJPs (reverse-mode AD), `"fwd_over_rev"` uses forward-over-reverse HVPs, `"rev_over_fwd"` uses reverse-over-forward HVPs, `"rev_over_rev"` uses reverse-over-reverse HVPs.

`save(path)` ¶

Save colored pattern to an .npz file.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination file path.	required

`load(path)` `classmethod` ¶

Load colored pattern from an .npz file.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source file path.	required

`repr()` ¶

Return compact single-line representation.

`str()` ¶

Render colored pattern with sparsity grid and color assignments.

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

AD mode for Jacobian computation.

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

`asdex.HessianMode = Literal['fwd_over_rev', 'rev_over_fwd', 'rev_over_rev']` `module-attribute` ¶

AD composition strategy for Hessian-vector products.

"fwd_over_rev" uses forward-over-reverse, "rev_over_fwd" uses reverse-over-forward, "rev_over_rev" uses reverse-over-reverse.

Full API¶

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

asdex.hessian(f, input_shape, *, mode=None, symmetric=True) ¶

asdex.jacobian_from_coloring(f, coloring) ¶

asdex.hessian_from_coloring(f, coloring) ¶

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

asdex.hessian_coloring(f, input_shape, *, mode=None, symmetric=True) ¶

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

asdex.hessian_coloring_from_sparsity(sparsity, *, mode=None, symmetric=True) ¶

asdex.jacobian_sparsity(f, input_shape) ¶

asdex.hessian_sparsity(f, input_shape) ¶

asdex.SparsityPattern dataclass ¶

nnz property ¶

m property ¶

n property ¶

density property ¶

col_to_rows cached property ¶

row_to_cols cached property ¶

__post_init__() ¶

from_coo(rows, cols, shape, *, input_shape=None) classmethod ¶

from_bcoo(bcoo) classmethod ¶

from_dense(dense) classmethod ¶

to_bcoo(data=None) ¶

todense() ¶

save(path) ¶

load(path) classmethod ¶

__str__() ¶

__repr__() ¶

asdex.ColoredPattern dataclass ¶

save(path) ¶

load(path) classmethod ¶

__repr__() ¶

__str__() ¶

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

asdex.HessianMode = Literal['fwd_over_rev', 'rev_over_fwd', 'rev_over_rev'] module-attribute ¶

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

`asdex.hessian(f, input_shape, *, mode=None, symmetric=True)` ¶

`asdex.jacobian_from_coloring(f, coloring)` ¶

`asdex.hessian_from_coloring(f, coloring)` ¶

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

`asdex.hessian_coloring(f, input_shape, *, mode=None, symmetric=True)` ¶

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

`asdex.hessian_coloring_from_sparsity(sparsity, *, mode=None, symmetric=True)` ¶

`asdex.jacobian_sparsity(f, input_shape)` ¶

`asdex.hessian_sparsity(f, input_shape)` ¶

`asdex.SparsityPattern` `dataclass` ¶

`nnz` `property` ¶

`m` `property` ¶

`n` `property` ¶

`density` `property` ¶

`col_to_rows` `cached` `property` ¶

`row_to_cols` `cached` `property` ¶

`__post_init__()` ¶

`from_coo(rows, cols, shape, *, input_shape=None)` `classmethod` ¶

`from_bcoo(bcoo)` `classmethod` ¶

`from_dense(dense)` `classmethod` ¶

`to_bcoo(data=None)` ¶

`todense()` ¶

`save(path)` ¶

`load(path)` `classmethod` ¶

`str()` ¶

`repr()` ¶

`asdex.ColoredPattern` `dataclass` ¶

`save(path)` ¶

`load(path)` `classmethod` ¶

`repr()` ¶

`str()` ¶

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

`asdex.HessianMode = Literal['fwd_over_rev', 'rev_over_fwd', 'rev_over_rev']` `module-attribute` ¶