Hessian

Hessian Computation

asdex.hessian(f, *args, argnums=0, has_aux=False, holomorphic=False, allow_int=False, mode=None, symmetric=True, output_format='bcoo')

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

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

asdex.value_and_hessian(f, *args, argnums=0, has_aux=False, holomorphic=False, allow_int=False, mode=None, symmetric=True, output_format='bcoo')

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

Like hessian, but also returns the primal value f(*args) without an extra forward pass.

asdex.hessian_from_coloring(f, coloring, output_format='bcoo', *, has_aux=False, holomorphic=False, allow_int=False)

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

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

asdex.value_and_hessian_from_coloring(f, coloring, output_format='bcoo', *, has_aux=False, holomorphic=False, allow_int=False)

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

Coloring

asdex.hessian_coloring(f, *args, argnums=0, has_aux=False, 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 one or more positional arrays.	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 (default 0).	0
has_aux	bool	If True, f is assumed to return (output, aux) where aux is auxiliary data ignored by sparsity detection.	False
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.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.

Sparsity Detection

asdex.hessian_sparsity(f, *args, argnums=0, has_aux=False)

Detect global Hessian sparsity pattern for a scalar-valued f.

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 one or more positional arrays.	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 (default 0).	0
has_aux	bool	Whether f returns (scalar_output, auxiliary_data). When True, aux is stripped before detection.	False

Returns:

Type	Description
SparsityPattern	Square SparsityPattern over the combined, selected input space.

Verification

asdex.check_hessian_correctness(f, x, coloring, *, method='matvec', num_probes=25, seed=0, rtol=None, atol=None)

Verify asdex's sparse Hessian against a JAX reference at a given input.

Parameters:

Name	Type	Description	Default
f	Callable[..., Any]	Scalar-valued function taking an array.	required
x	Any	Input at which to evaluate the Hessian. For multi-input functions (where argnums is a tuple), pass a tuple of all positional arguments.	required
coloring	ColoredPattern	Pre-computed colored pattern from :func:~asdex.hessian_coloring.	required
method	Literal['matvec', 'dense']	Verification method. "matvec" uses randomized matrix-vector products, which is O(k) in the number of probes. "dense" materializes the full dense Hessian, which is O(n^2).	'matvec'
num_probes	int	Number of random probe vectors (only used by "matvec").	25
seed	int	PRNG seed for reproducibility (only used by "matvec").	0
rtol	float | None	Relative tolerance for comparison. Defaults to 1e-5 for "matvec" and 1e-7 for "dense".	None
atol	float | None	Absolute tolerance for comparison. Defaults to 1e-5 for "matvec" and 1e-7 for "dense".	None

Raises:

Type	Description
VerificationError	If the sparse and reference Hessians disagree.

Configuration

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.