Skip to content

Full API

Differentiation

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

jacobian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: JacobianMode | None = None,
    symmetric: bool = False,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], BCOO]

jacobian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: JacobianMode | None = None,
    symmetric: bool = False,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], jax.Array]

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], BCOO | Array]

A function that takes an input array and returns the Jacobian of shape (m, n) where n = x.size and m = prod(output_shape). The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.value_and_jacobian(f, input_shape, *, mode=None, symmetric=False, output_format='bcoo') 

value_and_jacobian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: JacobianMode | None = None,
    symmetric: bool = False,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], tuple[jax.Array, BCOO]]

value_and_jacobian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: JacobianMode | None = None,
    symmetric: bool = False,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], tuple[jax.Array, jax.Array]]

Detect sparsity, color, and return a function computing value and sparse Jacobian.

Like jacobian, but also returns the primal value f(x) without an extra forward pass.

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], tuple[Array, BCOO | Array]]

A function that takes an input array and returns (f(x), J) where J is the Jacobian of shape (m, n) where n = x.size and m = prod(output_shape). The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.hessian(f, input_shape, *, mode=None, symmetric=True, output_format='bcoo') 

hessian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: HessianMode | None = None,
    symmetric: bool = True,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], BCOO]

hessian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: HessianMode | None = None,
    symmetric: bool = True,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], jax.Array]

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], BCOO | Array]

A function that takes an input array and returns the Hessian of shape (n, n) where n = x.size. The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.value_and_hessian(f, input_shape, *, mode=None, symmetric=True, output_format='bcoo') 

value_and_hessian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: HessianMode | None = None,
    symmetric: bool = True,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], tuple[jax.Array, BCOO]]

value_and_hessian(
    f: Callable[[ArrayLike], ArrayLike],
    input_shape: int | tuple[int, ...],
    *,
    mode: HessianMode | None = None,
    symmetric: bool = True,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], tuple[jax.Array, jax.Array]]

Detect sparsity, color, and return a function computing value and sparse Hessian.

Like hessian, but can also return the primal value f(x) without an extra forward pass.

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], tuple[Array, BCOO | Array]]

A function that takes an input array and returns (f(x), H) where H is the Hessian of shape (n, n) where n = x.size. The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.jacobian_from_coloring(f, coloring, output_format='bcoo') 

jacobian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], BCOO]

jacobian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], jax.Array]

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], BCOO | Array]

A function that takes an input array and returns the Jacobian of shape (m, n) where n = x.size and m = prod(output_shape). The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.value_and_jacobian_from_coloring(f, coloring, output_format='bcoo') 

value_and_jacobian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], tuple[jax.Array, BCOO]]

value_and_jacobian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], tuple[jax.Array, jax.Array]]

Build a function computing value and sparse Jacobian from a pre-computed coloring.

Like jacobian_from_coloring, but also returns the primal value f(x) without an extra forward pass.

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], tuple[Array, BCOO | Array]]

A function that takes an input array and returns (f(x), J) where J is the Jacobian of shape (m, n) where n = x.size and m = prod(output_shape). The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.hessian_from_coloring(f, coloring, output_format='bcoo') 

hessian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], BCOO]

hessian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], jax.Array]

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], BCOO | Array]

A function that takes an input array and returns the Hessian of shape (n, n) where n = x.size. The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

asdex.value_and_hessian_from_coloring(f, coloring, output_format='bcoo') 

value_and_hessian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["bcoo"] = ...,
) -> Callable[[ArrayLike], tuple[jax.Array, BCOO]]

value_and_hessian_from_coloring(
    f: Callable[[ArrayLike], ArrayLike],
    coloring: ColoredPattern,
    output_format: Literal["dense"],
) -> Callable[[ArrayLike], tuple[jax.Array, jax.Array]]

Build a function computing value and sparse Hessian from a pre-computed coloring.

Like hessian_from_coloring, but can also return the primal value f(x) without an extra forward pass.

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
output_format OutputFormat

Type of the output matrix. "bcoo" returns a sparse matrix of type jax.experimental.sparse.BCOO (default), "dense" returns a dense matrix of type jax.Array.

 'bcoo'

Returns:

Type Description
Callable[[ArrayLike], tuple[Array, BCOO | Array]]

A function that takes an input array and returns (f(x), H) where H is the Hessian of shape (n, n) where n = x.size. The output type depends on output_format: a sparse jax.experimental.sparse.BCOO (default) or a dense jax.Array.

Coloring

asdex.jacobian_coloring(f, input_shape, *, mode=None, symmetric=False, postprocess=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
postprocess bool

Only read when symmetric=True. Prune colors never used as hubs and compact the remaining ones (reduces the number of VJPs/JVPs during decompression). Defaults to False, matching SparseMatrixColorings.jl.

 False

Returns:

Type Description
ColoredPattern

A ColoredPattern ready for jacobian_from_coloring.

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

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
postprocess bool

Only read when symmetric=True. Prune colors never used as hubs and compact the remaining ones (reduces the number of HVPs during decompression). Defaults to False, matching SparseMatrixColorings.jl.

 False

Returns:

Type Description
ColoredPattern

A ColoredPattern ready for hessian_from_coloring.

asdex.jacobian_coloring_from_sparsity(sparsity, *, mode=None, symmetric=False, postprocess=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 | NDArray | BCOO

A SparsityPattern, NumPy array, or JAX BCOO matrix 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
postprocess bool

Only read when symmetric=True. Prune colors never used as hubs and compact the remaining ones (reduces the number of VJPs/JVPs during decompression). Defaults to False, matching SparseMatrixColorings.jl.

 False

Returns:

Type Description
ColoredPattern

A ColoredPattern ready for jacobian_from_coloring.

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

Color a sparsity pattern for sparse Hessian computation.

Parameters:

Name Type Description Default
sparsity SparsityPattern | NDArray | BCOO

A SparsityPattern, NumPy array, or JAX BCOO matrix 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
postprocess bool

Only read when symmetric=True. Prune colors never used as hubs and compact the remaining ones (reduces the number of HVPs during decompression). Pruned vertices get the neutral color -1 in the output (no HVP is computed for them). Defaults to False, matching SparseMatrixColorings.jl.

 False

Returns:

Type Description
ColoredPattern

A ColoredPattern ready for hessian_from_coloring.

asdex.color_rows(sparsity)

Greedy row-wise coloring for sparse Jacobian computation.

Assigns colors to rows such that no two rows sharing a non-zero column have the same color. This enables computing multiple Jacobian rows in a single VJP by using a combined seed vector.

Uses LargestFirst vertex ordering for fewer colors.

Parameters:

Name Type Description Default
sparsity SparsityPattern

SparsityPattern of shape (m, n) representing the Jacobian sparsity pattern

 required

Returns:

Type Description
tuple[NDArray[int32], int]

Tuple of (colors, num_colors) where:

  • colors: Array of shape (m,) with color assignment for each row
  • num_colors: Total number of colors used

asdex.color_cols(sparsity)

Greedy column-wise coloring for sparse Jacobian computation.

Assigns colors to columns such that no two columns sharing a non-zero row have the same color. This enables computing multiple Jacobian columns in a single JVP by using a combined tangent vector.

Uses LargestFirst vertex ordering for fewer colors.

Parameters:

Name Type Description Default
sparsity SparsityPattern

SparsityPattern of shape (m, n) representing the Jacobian sparsity pattern

 required

Returns:

Type Description
tuple[NDArray[int32], int]

Tuple of (colors, num_colors) where:

  • colors: Array of shape (n,) with color assignment for each column
  • num_colors: Total number of colors used

asdex.color_symmetric(sparsity, *, postprocess=False, forced_colors=None)

Greedy symmetric coloring for sparse Hessian computation.

Implements Algorithm 4.1 from Gebremedhin et al. (2007). A star coloring is a distance-1 coloring with the additional constraint that every path on 4 vertices uses at least 3 colors. Returns a :class:StarSet alongside the colors so that Hessian decompression can use hub-based extraction.

Uses LargestFirst vertex ordering.

Parameters:

Name Type Description Default
sparsity SparsityPattern

SparsityPattern of shape (n, n) representing the symmetric Hessian sparsity pattern.

 required
postprocess bool

If True, replace colors that are never used as a hub color (and not forced by a diagonal entry) with -1 (neutral), then compact remaining colors down. This reduces the number of HVPs needed during decompression. Defaults to False, matching SparseMatrixColorings.jl's postprocessing=false default.

 False
forced_colors NDArray[int32] | list[int] | None

Optional pre-computed color assignment of shape (n,). When provided, the algorithm verifies it satisfies the star-coloring constraints and raises :class:InvalidColoringError otherwise.

 None

Returns:

Type Description
tuple[NDArray[int32], int, StarSet]

Tuple (colors, num_colors, star_set) where:

  • colors: Array of shape (n,) with color assignment for each vertex. Values are in [0, num_colors - 1] for active vertices. After postprocessing, vertices whose color is pruned have value -1 (neutral — no HVP needed for them).
  • num_colors: Number of active colors (i.e. number of HVPs).
  • star_set: :class:StarSet encoding the 2-colored star decomposition.

Raises:

Type Description
ValueError

If pattern is not square.
InvalidColoringError

If forced_colors violates a star-coloring constraint.

Visualization

asdex.spy(pattern, *, ax=None, compressed=False, cmap=None, **kwargs)

Plot a sparsity pattern or colored pattern using matplotlib.

For a SparsityPattern, plots nonzeros as filled cells on a grid. For a ColoredPattern, fills cells with their assigned color.

When compressed=True on a ColoredPattern, plots the compressed pattern after coloring instead of the original.

Parameters:

Name Type Description Default
pattern SparsityPattern | ColoredPattern

The sparsity or colored pattern to plot.

 required
ax Axes | None

Matplotlib axes to plot on. If None, creates a new figure.

 None
compressed bool

If True and pattern is a ColoredPattern, plot the compressed pattern instead of the original.

 False
cmap Any

Matplotlib colormap for colored patterns. If None, uses tab10.

 None
**kwargs Any

Extra keyword arguments passed to ax.imshow.

 {}

Returns:

Type Description
Axes

The matplotlib axes with the plot.

Sparsity Detection

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.

Data Structures

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. A value of -1 means "neutral": the vertex is not seeded (used after star-coloring postprocessing).
num_colors int

Total number of active colors (number of JVPs/VJPs/HVPs).
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.
star_set StarSet | None

Star-coloring structure (hub/spoke assignment per edge). Present only for symmetric colorings produced by color_symmetric; None otherwise.

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.

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.