Skip to content

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. A matrix carrying indices_sorted adds the sparse_indices_sorted tag automatically, so a solver can skip its sort.

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. A matrix carrying indices_sorted adds the sparse_indices_sorted tag automatically, so a solver can skip its sort.

materialise_as_bcoo 

materialise_as_bcoo(
    operator: AbstractLinearOperator,
) -> BCOOLinearOperator

Materialise a sparse or tagged operator as a BCOOLinearOperator.

A BCOOLinearOperator is returned as it is, and a BCSRLinearOperator is converted. A lineax.JacobianLinearOperator or lineax.FunctionLinearOperator must carry a tag from splineax.sparsity_pattern_tag or splineax.sparsity_coloring_tag built from concrete indices. Its matrix is then computed with asdex, one JVP or VJP per color of the tag's coloring. The entries of the result are in the order of the tag's indices. A lineax.TaggedLinearOperator around either of them adds its own tags.

The direction of a tag's coloring is fixed if the tag holds one. Otherwise it is the jac of a lineax.JacobianLinearOperator, and asdex picks it when jac is None or for a lineax.FunctionLinearOperator.

Inputs and outputs may be pytrees, which are raveled in leaf order. Only real dtypes are supported.

sparsity_coloring_tag 

sparsity_coloring_tag(
    pattern: PatternSource,
    *,
    jac: JacobianDirection | None = None,
) -> _ContentPatternTag

sparsity_coloring_tag(
    pattern: Callable[..., Any],
    point: PyTree[
        Inexact[ArrayLike, ...] | ShapeDtypeStruct
    ],
    args: PyTree[Any] = None,
    *,
    jac: JacobianDirection | None = None,
) -> _ContentPatternTag

sparsity_coloring_tag(
    pattern: PatternSource | Callable[..., Any],
    point: PyTree[
        Inexact[ArrayLike, ...] | ShapeDtypeStruct
    ] = None,
    args: PyTree[Any] = None,
    *,
    jac: JacobianDirection | None = None,
) -> _ContentPatternTag

Create a sparsity-pattern tag that holds a Jacobian coloring.

A tagged lineax.JacobianLinearOperator or lineax.FunctionLinearOperator is turned into a BCOO with one JVP or VJP per color. This tag computes the coloring now, so no solver has to compute it later. The tag compares equal to a splineax.sparsity_pattern_tag of the same pattern, so operators with either tag will share a factorization.

Coloring runs host-side on numpy data, so call this outside jax.jit and close over the tag or pass it in as a static value.

Arguments:

  • pattern: either a known pattern, or a function fn(x, args) -> y whose Jacobian pattern is detected with asdex. A known pattern is a BCOO, a BCSR, a sparse operator from this package, an asdex.SparsityPattern, a dense boolean mask, or an asdex.ColoredPattern. A colored pattern is used as it is.
  • point: with a function, a point of the right structure, as arrays or jax.ShapeDtypeStructs. Only the shapes and dtypes matter. Inputs and outputs may be pytrees, which are raveled in leaf order.
  • args: with a function, extra arguments to it that are not differentiated.
  • jac: "fwd" colors columns for JVPs, "bwd" colors rows for VJPs. With None, asdex picks whichever needs fewer colors.

JacobianDirection module-attribute 

JacobianDirection = Literal['fwd', 'bwd']

The direction a Jacobian is computed in, named as in lineax.JacobianLinearOperator.

"fwd" colors columns and uses one JVP per color. "bwd" colors rows and uses one VJP per color.