Coloring

Sparsity Detection and Coloring

asdex.jacobian_coloring(f, *args, argnums=_DEFAULT_ARGNUMS, has_aux=_DEFAULT_HAS_AUX, mode=_DEFAULT_MODE, symmetric=_DEFAULT_SYMMETRIC_JACOBIAN, postprocess=_DEFAULT_POSTPROCESS, **kwargs)

Detect Jacobian sparsity and color in one step.

Parameters:

Name	Type	Description	Default
f	Callable	Function whose Jacobian is to be computed.	required
*args	Any	Sample arguments of f. Only structure and dtypes are used, values are ignored.	()
argnums	int | Sequence[int]	Specifies which positional argument(s) to differentiate with respect to. Defaults to 0.	_DEFAULT_ARGNUMS
has_aux	bool	Whether f returns (output, auxiliary_data). When True, the auxiliary output is ignored and only output is analyzed for sparsity.	_DEFAULT_HAS_AUX
mode	JacobianMode | None	AD mode. "fwd" uses JVPs (column coloring), "rev" uses VJPs (row coloring). Defaults to picking whichever of fwd/rev needs fewer colors (unless symmetric is True, in which case defaults to "fwd").	_DEFAULT_MODE
symmetric	bool	Whether to use symmetric coloring. Requires a square Jacobian. Defaults to False.	_DEFAULT_SYMMETRIC_JACOBIAN
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.	_DEFAULT_POSTPROCESS
**kwargs	Any	Sample keyword arguments of f. Non-traceable values (bools, strings, ints) are bound statically.	{}

Returns:

Type	Description
ColoredPattern	A ColoredPattern ready for jacobian_from_coloring.

asdex.hessian_coloring(f, *args, argnums=_DEFAULT_ARGNUMS, has_aux=_DEFAULT_HAS_AUX, mode=_DEFAULT_MODE, symmetric=_DEFAULT_SYMMETRIC_HESSIAN, postprocess=_DEFAULT_POSTPROCESS, **kwargs)

Detect Hessian sparsity and color in one step.

Parameters:

Name	Type	Description	Default
f	Callable	Scalar-valued function whose Hessian is to be computed.	required
*args	Any	Sample arguments of f. Only structure and dtypes are used, values are ignored.	()
argnums	int | Sequence[int]	Specifies which positional argument(s) to differentiate with respect to. Defaults to 0.	_DEFAULT_ARGNUMS
has_aux	bool	Whether f returns (output, auxiliary_data). When True, the auxiliary output is ignored and only output is analyzed for sparsity.	_DEFAULT_HAS_AUX
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".	_DEFAULT_MODE
symmetric	bool	Whether to use symmetric coloring. Defaults to True.	_DEFAULT_SYMMETRIC_HESSIAN
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 (no HVP is computed for them). Defaults to False.	_DEFAULT_POSTPROCESS
**kwargs	Any	Sample keyword arguments of f. Non-traceable values (bools, strings, ints) are bound statically.	{}

Returns:

Type	Description
ColoredPattern	A ColoredPattern ready for hessian_from_coloring.

Coloring Known Sparsity Patterns

asdex.jacobian_coloring_from_sparsity(sparsity, *, mode=_DEFAULT_MODE, symmetric=_DEFAULT_SYMMETRIC_JACOBIAN, postprocess=_DEFAULT_POSTPROCESS)

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). Defaults to picking whichever of fwd/rev needs fewer colors (unless symmetric is True, in which case defaults to "fwd").	_DEFAULT_MODE
symmetric	bool	Whether to use symmetric coloring. Requires a square Jacobian. Defaults to False.	_DEFAULT_SYMMETRIC_JACOBIAN
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.	_DEFAULT_POSTPROCESS

Returns:

Type	Description
ColoredPattern	A ColoredPattern ready for jacobian_from_coloring.

asdex.hessian_coloring_from_sparsity(sparsity, *, mode=_DEFAULT_MODE, symmetric=_DEFAULT_SYMMETRIC_HESSIAN, postprocess=_DEFAULT_POSTPROCESS)

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".	_DEFAULT_MODE
symmetric	bool	Whether to use symmetric coloring. Defaults to True.	_DEFAULT_SYMMETRIC_HESSIAN
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 (no HVP is computed for them). Defaults to False.	_DEFAULT_POSTPROCESS

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	Jacobian sparsity pattern of shape (m, n).	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	Jacobian sparsity pattern of shape (m, n).	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=_DEFAULT_POSTPROCESS, 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	Hessian sparsity pattern of shape (n, n).	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.	_DEFAULT_POSTPROCESS
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 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.

Validation

asdex.check_coloring_rows(sparsity, colors)

Check a row coloring: no column contains two rows with the same color.

Parameters:

Name	Type	Description	Default
sparsity	SparsityPattern	Jacobian sparsity pattern of shape (m, n).	required
colors	NDArray[int32]	Row color assignment, shape (m,).	required

Raises:

Type	Description
InvalidColoringError	If the coloring is invalid.

asdex.check_coloring_cols(sparsity, colors)

Check a column coloring: no row contains two columns with the same color.

Parameters:

Name	Type	Description	Default
sparsity	SparsityPattern	Jacobian sparsity pattern of shape (m, n).	required
colors	NDArray[int32]	Column color assignment, shape (n,).	required

Raises:

Type	Description
InvalidColoringError	If the coloring is invalid.

asdex.check_coloring_symmetric(sparsity, colors)

Check a star coloring: distance-1 coloring + no 2-colored path of 4 vertices.

A star coloring satisfies:

1. Adjacent vertices have different colors (distance-1).
2. Every path on 4 vertices uses at least 3 distinct colors (no 2-colored P4).

Vertices with color -1 (neutral, from postprocessing) are excluded from the P4 check as they represent the absence of a color.

Parameters:

Name	Type	Description	Default
sparsity	SparsityPattern	Hessian sparsity pattern of shape (n, n).	required
colors	NDArray[int32]	Vertex color assignment, shape (n,).	required

Raises:

Type	Description
ValueError	If pattern is not square.
InvalidColoringError	If the coloring is invalid.

asdex.InvalidColoringError

Bases: ValueError

Raised when a user-supplied coloring violates a star-coloring constraint.

See color_symmetric with forced_colors.