jax_sn.operators.base_linear_operator module

This module defines specialized base classes for linear operators in jax-sn, using the lineax operators as a base.

The strategy + data is used here. Data is a pytree that contains all data the strategy needs to apply the operator.

This cleanly separates choice (i.e. batch size) from the rest of the data. All operators are then of similar nature, with some specialized classes for common patterns (e.g. transpose boolean vs data transpose).

Why not inheriting and implementing a forward?

This works well enough for single operators, but once it starts composing, it gets messy: if you want to have implementations of a different kind, you need to have all the data for such an operator in the structure, leading to duplicated data (not one source of truth).

Consider:

class Bop:
    fem_data : FEMData

    @abstractmethod
    def mv(self, x):
        ...

class Cop:
    fem_data : FEMData

    @abstractmethod
    def mv(self, x):
        ...

class Aop:
    bop : Bop
    cop : Cop

    def mv(self, x):
        return self.bop.mv(self.cop.mv(x))

The data for Bop and Cop is now duplicated twice! Instead we have:

class Bop:
    fem_data : FEMData
    strategy : Bstrategy

    def mv(self, x):
        return self.strategy.forward(self.fem_data, x)

class Cop:
    fem_data : FEMData
    strategy : Cstrategy

    def mv(self, x):
        return self.strategy.forward(self.fem_data, x)

class ComposedA(Astrategy):
    b_strategy : Bstrategy
    c_strategy : Cstrategy

    def forward(self, fem_data, x):
        return self.b_strategy.forward(fem_data, self.c_strategy.forward(fem_data, x))

class Aop:
    fem_data   : FEMData
    a_strategy : Astrategy

    def mv(self, x):
        return self.a_strategy.forward(self.fem_data, x)

Construction is free (it’s just composing, no computation), and the data is only stored once. The strategies are stateless, so they can be shared between operators. This further gives us the flexibility to have fused strategies, which make no reference to their suboperators (i.e. a multi-angle fused sweep!)

class AbstractLinearOperator_sn

Bases: AbstractLinearOperator

Base class for linear operators in jax-sn. Extends lineax AbstractLinearOperator to provide a common base class for all jax-sn linear operators.

Some specific functionality for jax-sn linear operators can be added here in the future.

Since the operators can have very large sizes, we prefer to throw by default when as_matrix is called. Since we only use real operators in jax-sn, the conj method is implemented to return self.

Abstract methods are forwarded from lineax.AbstractLinearOperator, and must be implemented by subclasses.

as_matrix()

In jax-sn, as_matrix is not implemented for linear operators due to potentially extremely large size. Prefer throwing to silently materializing.

Return type:

Inexact[Array, 'a b']

conj()

All operators in jax-sn are real, so the conjugate is the operator itself.

Return type:

Self

abstractmethod mv(x: PyTree[jaxtyping.Inexact[Array, '_b']])

Apply the operator to a vector.

Return type:

Inexact[Array, '_a']]

abstractmethod transpose()

Return the transposed operator.

Return type:

AbstractLinearOperator_sn

abstractmethod in_structure()

Return the structure of the input vector.

Return type:

ShapeDtypeStruct]

abstractmethod out_structure()

Return the structure of the output vector.

Return type:

ShapeDtypeStruct]

class DirectlyInvertibleOperator

Bases: AbstractLinearOperator_sn

An operator that can be directly inverted. Uses lineax to define the JVP and VJP.

inv(b: PyTree[jaxtyping.Inexact[Array, '_a']])

Solve \(\mathbf{A} x = b\) for x using the operator’s own direct inverse.

Return type:

Inexact[Array, '_b']]

class DirectInverseSolver

Bases: AbstractLinearSolver[DirectlyInvertibleOperator]

Complete dummy class for always-correct linear solves.

init(operator: AbstractLinearOperator, options: dict[str, Any])

Does nothing, but checks that the operator is a DirectlyInvertibleOperator.

Return type:

DirectlyInvertibleOperator

compute(state: DirectlyInvertibleOperator, vector: PyTree[jax.Array], options: dict[str, Any])

Always a succesful solve using the internal inversion

Return type:

tuple[Array], RESULTS, dict[str, Any]]

assume_full_rank()

Always full rank

Return type:

bool

transpose(state: DirectlyInvertibleOperator, options: dict[str, Any])

Transpose the solver state or equally the operator

Return type:

tuple[DirectlyInvertibleOperator, dict[str, Any]]

conj(state: DirectlyInvertibleOperator, options: dict[str, Any])

Conj defers to the operator.

Return type:

tuple[DirectlyInvertibleOperator, dict[str, Any]]

class BaseOperatorData

Bases: Module

Base data class. For type checking only (all can just be eqx.Module themselves).

class TransposableOperatorData

Bases: BaseOperatorData

Data that carries the direction of its operator, and so knows its own transpose (e.g. the transport data, whose transpose negates the angle and reverses the sweep ordering together).

abstract property transposed: Self

The data of the transposed operator.

class Strategy

Bases: Module

Applies a linear operator to a vector. Not transposable.

abstractmethod forward(data: BaseOperatorData, x: PyTree[jaxtyping.Inexact[Array, '_b']])

Return type:

Inexact[Array, '_a']]

class TransposableStrategy

Bases: Strategy

A strategy that implements the transposed action itself, for operators that carry the direction as a flag rather than in their data.

abstractmethod transpose(data: BaseOperatorData, x: PyTree[jaxtyping.Inexact[Array, '_a']])

Return type:

Inexact[Array, '_b']]

class SymmetricStrategy

Bases: TransposableStrategy

A strategy which is symmetric, i.e. forward and transpose are the same.

transpose(data: BaseOperatorData, x: PyTree[jaxtyping.Inexact[Array, '_a']])

Return type:

Inexact[Array, '_b']]

class DataStrategyLinearOperator

Bases: AbstractLinearOperator_sn

A linear operator which is defined by a data object and a strategy object. The data object contains the information needed to construct the operator, and the strategy object defines how to apply the operator.

data: AbstractVar[BaseOperatorData]

strategy: AbstractVar[Strategy]

abstractmethod mv(x: PyTree[jaxtyping.Inexact[Array, '_b']])

Apply the operator to a vector.

Return type:

Inexact[Array, '_a']]

class ForwardDataStrategyLinearOperator

Bases: DataStrategyLinearOperator

mv is always the strategy’s forward action, and in/out structures coincide. Subclasses supply the shared structure via _structure, and say how transposition works.

mv(x: PyTree[jaxtyping.Inexact[Array, '_b']])

Apply the operator to a vector.

Return type:

Inexact[Array, '_a']]

in_structure()

Return the structure of the input vector.

Return type:

ShapeDtypeStruct]

out_structure()

Return the structure of the output vector.

Return type:

ShapeDtypeStruct]

class SymmetricDataStrategyLinearOperator

Bases: ForwardDataStrategyLinearOperator

Transposition is a no-op: the strategy’s forward action is assumed to be the symmetric one.

transpose()

Return the transposed operator.

Return type:

Self

class DataTransposeDataStrategyLinearOperator

Bases: ForwardDataStrategyLinearOperator

Transposition is a transform of the data.

data: AbstractVar[TransposableOperatorData]

transpose()

Return the transposed operator.

Return type:

Self

class TransposeBoolDataStrategyLinearOperator

Bases: DataStrategyLinearOperator

Transposition is a flag: the data is direction-agnostic, so the strategy implements both directions and mv picks one. In/out structures swap with it.

strategy: AbstractVar[TransposableStrategy]

transpose_bool: AbstractVar[bool]

mv(x: PyTree[jaxtyping.Inexact[Array, '_b']])

Apply the operator to a vector.

Return type:

Inexact[Array, '_a']]

transpose()

Return the transposed operator.

Return type:

Self

in_structure()

Return the structure of the input vector.

Return type:

ShapeDtypeStruct]

out_structure()

Return the structure of the output vector.

Return type:

ShapeDtypeStruct]