Coloring

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.