Operators

BCOOLinearOperator

BCOOLinearOperator(
    matrix: Inexact[BCOO, "a b"],
    tags: object | frozenset[object] = (),
)

Bases: AbstractLinearOperator

Wraps a jax.experimental.sparse.BCOO array into a linear operator.

If the matrix has shape (a, b) then matrix-vector multiplication (self.mv) is defined in the usual way: as accepting a vector of shape (b,) and returning a vector of shape (a,).

Arguments:

• matrix: a two-dimensional BCOO array. For an array with shape (a, b) then this operator can perform matrix-vector products on a vector of shape (b,) to return a vector of shape (a,).
• tags: any tags indicating whether this matrix has any particular properties, like symmetry or positive-definite-ness. Note that these properties are unchecked and you may get incorrect values elsewhere if these tags are wrong.

---

BCSRLinearOperator

BCSRLinearOperator(
    matrix: Inexact[BCSR, "a b"],
    tags: object | frozenset[object] = (),
)

Bases: AbstractLinearOperator

Wraps a jax.experimental.sparse.BCSR array into a linear operator.

If the matrix has shape (a, b) then matrix-vector multiplication (self.mv) is defined in the usual way: as accepting a vector of shape (b,) and returning a vector of shape (a,).

Arguments:

• matrix: a two-dimensional BCSR array. For an array with shape (a, b) then this operator can perform matrix-vector products on a vector of shape (b,) to return a vector of shape (a,).
• tags: any tags indicating whether this matrix has any particular properties, like symmetry or positive-definite-ness. Note that these properties are unchecked and you may get incorrect values elsewhere if these tags are wrong.

---

SparseJacobianLinearOperator

SparseJacobianLinearOperator(
    fn: Callable,
    x: Inexact[ArrayLike, " n"],
    args: PyTree[Any] = None,
    *,
    sparsity: SparsityPattern
    | ndarray
    | BCOO
    | None = None,
    coloring: ColoredPattern
    | JacobianColoring
    | None = None,
    mode: JacobianMode | None = None,
    tags: object | frozenset[object] = (),
    transposed: bool = False,
    closure_convert: bool = True,
)

Bases: AbstractLinearOperator

Given a function fn: X -> Y and a point x, the linear operator (d(fn)/dx)(x), kept sparse.

The Jacobian's sparsity pattern and a matching coloring are determined once at construction (via asdex), so materialising the Jacobian at x costs one JVP or VJP per color rather than one per column or row. Materialise it with as_bcoo or lineax.materialise (sparse) or as_matrix (dense), or hand the operator straight to a splineax sparse solver, which materialises it through lineax.materialise.

The coloring is stored as an asdex.ColoredPattern, which is a registered JAX pytree, so the operator carries it as an ordinary (dynamic) field and the whole operator can be passed as an argument into a jitted function. A precomputed coloring may be supplied through the coloring argument (either an asdex.ColoredPattern or a splineax.JacobianColoring) to skip detection, which is what makes it cheap to build many operators for the same sparsity.

Only real dtypes are supported. To build many operators for the same function at different points without repeating sparsity detection, use splineax.SparseJacobianLinearOperatorColoring.

Arguments:

• fn: a function (x, args) -> y, where both x and y are one-dimensional arrays of real dtype. Its Jacobian d(fn)/dx is the linear operator.
• x: the point at which to evaluate d(fn)/dx.
• args: extra arguments to fn that are not differentiated.
• sparsity: optional known sparsity pattern of the Jacobian, as an asdex.SparsityPattern, a dense boolean mask, or a BCOO matrix. Skips sparsity detection (the pattern is still colored here).
• coloring: optional precomputed coloring, either an asdex.ColoredPattern or a splineax.JacobianColoring. Skips both sparsity detection and coloring. At most one of sparsity and coloring may be given.
• mode: optional asdex coloring mode, either "fwd" (column coloring, materialised with JVPs) or "rev" (row coloring, materialised with VJPs). If not given, asdex picks based on the pattern.
• tags: any lineax tags indicating whether the Jacobian has any particular properties, like symmetry or positive-definite-ness. Note that these properties are unchecked and you may get incorrect values elsewhere if these tags are wrong.

transposed and closure_convert are internal arguments, used by transpose() and splineax.SparseJacobianLinearOperatorColoring.

as_bcoo

as_bcoo() -> Inexact[BCOO, 'a b']

Materialises the Jacobian at x as a BCOO matrix, using one JVP or VJP per color of the precomputed coloring.

---

JacobianColoring

Bases: Module

A function-agnostic Jacobian sparsity coloring, backed by asdex.

A JacobianColoring wraps an asdex.ColoredPattern: a Jacobian sparsity pattern together with a matching row or column coloring. The coloring is what lets the Jacobian be materialised with one JVP or VJP per color rather than one per column or row.

This wrapper carries only the coloring, not any particular function or evaluation point. Create one with either splineax.JacobianColoring.detect, which detects the sparsity of a function and colors it, or splineax.JacobianColoring.from_sparsity, which colors a sparsity pattern you already know. Both run host-side (they use numpy, scipy, and graph coloring under the hood), so build the coloring outside jax.jit and pass the finished object in as an argument.

Because asdex.ColoredPattern is a registered JAX pytree, a JacobianColoring is itself a pytree and can be threaded through jitted computations. Any two colorings of the same sparsity pattern flatten to the same treedef, so a jitted function that receives a JacobianColoring compiles once and stays cached even when the coloring is regenerated from scratch.

To turn a coloring into a linear operator, either pass it as the coloring argument of splineax.SparseJacobianLinearOperator together with a function and a point, or bind it to a specific function once with splineax.SparseJacobianLinearOperatorColoring.from_jacobian_coloring and then call operator_at at many points.

coloring instance-attribute

coloring: ColoredPattern

The wrapped asdex coloring, holding both the sparsity pattern and the row or column coloring of it. Stored as an ordinary (dynamic) pytree field, so it can be carried through jitted functions.

sparsity property

sparsity: SparsityPattern

The asdex.SparsityPattern that was colored. The splineax solvers read the row and column indices from here to pre-analyze the sparsity host-side.

mode property

mode: str

The resolved asdex coloring mode, either "fwd" or "rev". This is the mode asdex chose, never the unresolved None the caller may have passed.

num_colors property

num_colors: int

The number of colors, and so the number of JVPs or VJPs one Jacobian materialisation costs.

detect classmethod

detect(
    fn: Callable,
    x: Inexact[ArrayLike, " n"] | ShapeDtypeStruct,
    args: PyTree[Any] = None,
    *,
    mode: JacobianMode | None = None,
) -> JacobianColoring

Detects the Jacobian sparsity of fn and colors it.

Detection is structural: asdex traces the computation graph of fn and reads off which outputs depend on which inputs, without evaluating any derivatives. Only the shape and dtype of x matter, so a jax.ShapeDtypeStruct may be passed in place of a concrete point. Detection and coloring run host-side, so call this outside jax.jit and pass the resulting coloring in.

Arguments:

• fn: a function (x, args) -> y, where both x and y are one-dimensional arrays of real dtype. Its Jacobian's sparsity is detected.
• x: a representative point, or a jax.ShapeDtypeStruct describing one. Only its shape and dtype are used.
• args: extra arguments to fn that are not differentiated.
• mode: optional asdex coloring mode, either "fwd" (column coloring, materialised with JVPs) or "rev" (row coloring, materialised with VJPs). If not given, asdex picks based on the pattern.

from_sparsity classmethod

from_sparsity(
    sparsity: SparsityPattern | ndarray | BCOO,
    *,
    mode: JacobianMode | None = None,
) -> JacobianColoring

Colors a known Jacobian sparsity pattern, skipping detection.

No function is needed, since the sparsity pattern already describes which Jacobian entries are nonzero. Coloring runs host-side, so call this outside jax.jit and pass the resulting coloring in.

Arguments:

• sparsity: the known sparsity pattern of the Jacobian, as an asdex.SparsityPattern, a dense boolean mask, or a BCOO matrix.
• mode: optional asdex coloring mode, either "fwd" (column coloring, materialised with JVPs) or "rev" (row coloring, materialised with VJPs). If not given, asdex picks based on the pattern.

---

SparseJacobianLinearOperatorColoring

Bases: Module

A splineax.JacobianColoring bound to a specific function, reusable across evaluation points.

Where a splineax.JacobianColoring carries only the coloring, this class also holds the (closure-converted) function whose Jacobian was colored. That pairing is what a splineax.SparseJacobianLinearOperator needs, so splineax.SparseJacobianLinearOperatorColoring.operator_at can produce an operator at any point without repeating sparsity detection or coloring.

Build one with splineax.SparseJacobianLinearOperatorColoring.detect or splineax.SparseJacobianLinearOperatorColoring.from_sparsity, or from an existing splineax.JacobianColoring with splineax.SparseJacobianLinearOperatorColoring.from_jacobian_coloring. All operators built from one instance share the same closure-converted function and the same coloring pattern, so passing them through the same jitted computation compiles only once.

The coloring is valid for any x and args of the same abstract structure (shapes and dtypes) as those it was computed with. asdex guarantees that the sparsity pattern depends only on the traced computation graph, not on numerical values, so reusing the coloring at other points is always sound.

fn instance-attribute

fn: Callable

The closure-converted function whose Jacobian was colored. Shared by every operator built through operator_at, so their pytree structures compare equal.

coloring instance-attribute

coloring: JacobianColoring

The function-agnostic coloring. Carried as an ordinary (dynamic) pytree field, since JacobianColoring wraps a pytree asdex.ColoredPattern.

from_jacobian_coloring classmethod

from_jacobian_coloring(
    coloring: JacobianColoring,
    fn: Callable,
    x: Inexact[ArrayLike, " n"] | ShapeDtypeStruct,
    args: PyTree[Any] = None,
) -> SparseJacobianLinearOperatorColoring

Binds an existing splineax.JacobianColoring to a function.

This is the bridge from a bare coloring to an operator factory. The coloring may have come from splineax.JacobianColoring.detect on this same function or from splineax.JacobianColoring.from_sparsity on a pattern you know matches fn. The function is closure-converted here (once), and the result reused by every operator_at call.

Arguments:

• coloring: the coloring to bind, as a splineax.JacobianColoring.
• fn: a function (x, args) -> y, where both x and y are one-dimensional arrays of real dtype. Its Jacobian must have the sparsity the coloring describes.
• x: a representative point, or a jax.ShapeDtypeStruct describing one. Only its shape and dtype matter here, used to closure-convert fn.
• args: extra arguments to fn that are not differentiated.

detect classmethod

detect(
    fn: Callable,
    x: Inexact[ArrayLike, " n"] | ShapeDtypeStruct,
    args: PyTree[Any] = None,
    *,
    mode: JacobianMode | None = None,
) -> SparseJacobianLinearOperatorColoring

Detects the Jacobian sparsity of fn, colors it, and binds it to fn.

A convenience wrapper equivalent to splineax.JacobianColoring.detect followed by splineax.SparseJacobianLinearOperatorColoring.from_jacobian_coloring.

Arguments:

• fn: a function (x, args) -> y, where both x and y are one-dimensional arrays of real dtype.
• x: a representative point, or a jax.ShapeDtypeStruct describing one. Only its shape and dtype matter, since sparsity detection is structural.
• args: extra arguments to fn that are not differentiated.
• mode: optional asdex coloring mode, "fwd" or "rev".

from_sparsity classmethod

from_sparsity(
    fn: Callable,
    x: Inexact[ArrayLike, " n"] | ShapeDtypeStruct,
    sparsity: SparsityPattern | ndarray | BCOO,
    args: PyTree[Any] = None,
    *,
    mode: JacobianMode | None = None,
) -> SparseJacobianLinearOperatorColoring

Colors a known Jacobian sparsity pattern of fn, skipping detection, and binds it to fn.

A convenience wrapper equivalent to splineax.JacobianColoring.from_sparsity followed by splineax.SparseJacobianLinearOperatorColoring.from_jacobian_coloring.

Arguments:

• fn: a function (x, args) -> y, where both x and y are one-dimensional arrays of real dtype.
• x: a representative point, or a jax.ShapeDtypeStruct describing one. Not used for detection, but required to closure-convert fn.
• sparsity: the known sparsity pattern of the Jacobian, as an asdex.SparsityPattern, a dense boolean mask, or a BCOO matrix.
• args: extra arguments to fn that are not differentiated.
• mode: optional asdex coloring mode, "fwd" or "rev".

operator_at

operator_at(
    x: Inexact[ArrayLike, " n"],
    args: PyTree[Any] = None,
    tags: object | frozenset[object] = (),
) -> SparseJacobianLinearOperator

Builds a splineax.SparseJacobianLinearOperator at the point x, reusing the precomputed coloring.

Arguments:

• x: the point at which to evaluate the Jacobian. Must have the same shape and dtype as the point the coloring was computed for.
• args: extra arguments to fn that are not differentiated. Must have the same abstract structure as the args the coloring was computed for.
• tags: any lineax tags for the resulting operator.