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.
|
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 |
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 |
required |
Returns:
| Type | Description |
|---|---|
Callable[[ArrayLike], BCOO]
|
A function that takes an input array and returns
the sparse Hessian as BCOO of shape |
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.
|
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 |
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.
|
None
|
symmetric
|
bool
|
Whether to use symmetric (star) coloring. Defaults to True (exploits Hessian symmetry for fewer colors). |
True
|
Returns:
| Type | Description |
|---|---|
ColoredPattern
|
A |
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 |
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: |
required |
method
|
Literal['matvec', 'dense']
|
Verification method.
|
'matvec'
|
num_probes
|
int
|
Number of random probe vectors (only used by |
25
|
seed
|
int
|
PRNG seed for reproducibility (only used by |
0
|
rtol
|
float | None
|
Relative tolerance for comparison.
Defaults to 1e-5 for |
None
|
atol
|
float | None
|
Absolute tolerance for comparison.
Defaults to 1e-5 for |
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.