The idea is to have things like matmuls (and other opts) as dataclass attributes: ```python @dataclass class LinearOperator(...): ... matmul: Callable[[LinearOperator, Float[Array, "N M"]] Float[Array, "N M"]] = static_field(...) def __matmul__(self, other : Float[Array, "N M"]) -> Float[Array, "N M"] return matmul(self, other) ``` So that the user, can, create a `LinearOperator` from defining a matmul for conjugate gradients (akin to distrax's Lambda Bijector): ``` MyLinOp = LinearOperator(..., matmul=*) ``` This is clean to do from the recent refactoring of #20.
The idea is to have things like matmuls (and other opts) as dataclass attributes:
So that the user, can, create a
LinearOperatorfrom defining a matmul for conjugate gradients (akin to distrax's Lambda Bijector):This is clean to do from the recent refactoring of #20.