Coloring

Sparsity Detection and 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.

Coloring Known Sparsity Patterns

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.

Low-Level Algorithms

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.