howsiyu

from functools import partial
from typing import Any, NamedTuple, Protocol
import jax
import jax.numpy as jnp
from jax import lax
from jax import Array


class LinearOperator(Protocol):
    dtype: Any