Skip to content

Hessian

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

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

Returns:

Type Description
Callable[[ArrayLike], BCOO]

A function that takes an input array and returns the sparse Hessian as BCOO of shape (n, n) where n = x.size.

asdex.hessian_from_coloring(f, coloring)

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

Returns:

Type Description
Callable[[ArrayLike], BCOO]

A function that takes an input array and returns the sparse Hessian as BCOO of shape (n, n) where n = x.size.

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

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

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.

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[[ArrayLike], ArrayLike]

Scalar-valued function taking an array.

required
x ArrayLike

Input at which to evaluate the Hessian.

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.


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.