Skip to content
This repository was archived by the owner on May 11, 2023. It is now read-only.
This repository was archived by the owner on May 11, 2023. It is now read-only.

dev: Custom matmul, and LinOp ops constructer. #23

@daniel-dodd

Description

@daniel-dodd

The idea is to have things like matmuls (and other opts) as dataclass attributes:

@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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions