jaxonlayers 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

{jaxonlayers-0.1.0.dist-info → jaxonlayers-0.1.1.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: jaxonlayers
-Version: 0.1.0
+Version: 0.1.1
 Summary: Add your description here
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 Requires-Dist: beartype>=0.21.0
 Requires-Dist: equinox>=0.13.0
-Requires-Dist: jax>=0.7.2
+Requires-Dist: jax>=0.8.0
 Requires-Dist: jaxtyping>=0.3.2

jaxonlayers-0.1.1.dist-info/RECORD

@@ -0,0 +1,5 @@
+jaxonlayers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jaxonlayers-0.1.1.dist-info/METADATA,sha256=aLnCVoHFmkHYSU5QiDRhYhuYBN4IcZDXAmpuYm-asfk,275
+jaxonlayers-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jaxonlayers-0.1.1.dist-info/top_level.txt,sha256=n5UHFDErh3dJY77ypkEKlwFOQffKBnpGH9nUrwCinto,12
+jaxonlayers-0.1.1.dist-info/RECORD,,

jaxonlayers/functions/__init__.py

@@ -1,27 +0,0 @@
-from .attention import multi_head_attention_forward, shifted_window_attention
-from .initialization import kaiming_init_conv2d
-from .masking import (
-    build_attention_mask,
-    canonical_attn_mask,
-    canonical_key_padding_mask,
-    canonical_mask,
-)
-from .regularization import dropout, stochastic_depth
-from .state_space import selective_scan
-from .utils import (
-    default_floating_dtype,
-)
-
-__all__ = [
-    "multi_head_attention_forward",
-    "kaiming_init_conv2d",
-    "build_attention_mask",
-    "canonical_attn_mask",
-    "canonical_key_padding_mask",
-    "canonical_mask",
-    "stochastic_depth",
-    "selective_scan",
-    "dropout",
-    "default_floating_dtype",
-    "shifted_window_attention",
-]

jaxonlayers/functions/attention.py

@@ -1,374 +0,0 @@
-import functools
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-from beartype.typing import Any
-from jaxtyping import Array, Bool, Float, PRNGKeyArray
-
-from jaxonlayers.functions.normalization import normalize
-from jaxonlayers.functions.regularization import dropout as dropout_fn
-from jaxonlayers.functions.utils import default_floating_dtype
-
-
-def multi_head_attention_forward(
-    query: Float[Array, "tgt_len d_model"],
-    key: Float[Array, "src_len d_model"],
-    value: Float[Array, "src_len d_model"],
-    embed_dim_to_check: int,
-    num_heads: int,
-    in_proj_weight: Float[Array, "3*d_model d_model"] | None = None,
-    in_proj_bias: Float[Array, "3*d_model"] | None = None,
-    bias_k: Float[Array, "1 d_model"] | None = None,
-    bias_v: Float[Array, "1 d_model"] | None = None,
-    add_zero_attn: bool = False,
-    dropout_p: float = 0.0,
-    out_proj_weight: Float[Array, "d_model d_model"] | None = None,
-    out_proj_bias: Float[Array, "d_model"] | None = None,
-    inference: bool = False,
-    key_padding_mask: Float[Array, "src_len"] | Bool[Array, "src_len"] | None = None,
-    attn_mask: Float[Array, "tgt_len src_len"] | None = None,
-    need_weights: bool = True,
-    use_separate_proj_weight: bool = False,
-    q_proj_weight: Float[Array, "d_model d_model"] | None = None,
-    k_proj_weight: Float[Array, "d_model d_model"] | None = None,
-    v_proj_weight: Float[Array, "d_model d_model"] | None = None,
-    static_k: Float[Array, "src_len d_model"] | None = None,
-    static_v: Float[Array, "src_len d_model"] | None = None,
-    average_attn_weights: bool = True,
-    is_causal: bool = False,
-    dropout_key: PRNGKeyArray | None = None,
-) -> tuple[
-    Float[Array, "tgt_len d_model"],
-    Float[Array, "num_heads tgt_len src_len"]
-    | Float[Array, "tgt_len src_len"]
-    | Float[Array, "tgt_len src_len+1"]
-    | None,
-]:
-    tgt_len, d_model = query.shape
-    src_len, k_dim = key.shape
-    value_len, v_dim = value.shape
-
-    assert d_model == k_dim == v_dim == embed_dim_to_check, (
-        "Embedding dimensions must match"
-    )
-
-    assert src_len == value_len, "Key and value must have the same sequence length"
-
-    head_dim = d_model // num_heads
-    assert head_dim * num_heads == d_model, "embed_dim must be divisible by num_heads"
-
-    if dropout_p > 0.0:
-        assert dropout_key is not None, (
-            "dropout_key must be provided if dropout_p > 0.0"
-        )
-
-    if use_separate_proj_weight:
-        # When using separate projection weights for q, k, v
-        assert q_proj_weight is not None, (
-            "q_proj_weight should not be None when use_separate_proj_weight=True"
-        )
-        assert k_proj_weight is not None, (
-            "k_proj_weight should not be None when use_separate_proj_weight=True"
-        )
-        assert v_proj_weight is not None, (
-            "v_proj_weight should not be None when use_separate_proj_weight=True"
-        )
-
-        q = query @ q_proj_weight.T
-
-        if static_k is None:
-            k = key @ k_proj_weight.T
-        else:
-            k = static_k
-            src_len, _ = k.shape
-
-        if static_v is None:
-            v = value @ v_proj_weight.T
-        else:
-            v = static_v
-            value_len, _ = v.shape
-
-        if in_proj_bias is not None:
-            q_bias, k_bias, v_bias = jnp.split(in_proj_bias, 3)
-            q = q + q_bias
-            k = k + k_bias
-            v = v + v_bias
-
-    else:
-        assert in_proj_weight is not None, (
-            "in_proj_weight should not be None when use_separate_proj_weight=False"
-        )
-
-        q_proj_weight_part, k_proj_weight_part, v_proj_weight_part = jnp.split(
-            in_proj_weight, 3
-        )
-
-        q = query @ q_proj_weight_part.T
-
-        if static_k is None:
-            k = key @ k_proj_weight_part.T
-        else:
-            k = static_k
-            src_len, _ = static_k.shape
-
-        if static_v is None:
-            v = value @ v_proj_weight_part.T
-        else:
-            v = static_v
-            value_len, _ = static_v.shape
-
-        if in_proj_bias is not None:
-            q_bias, k_bias, v_bias = jnp.split(in_proj_bias, 3)
-            q = q + q_bias
-            k = k + k_bias
-            v = v + v_bias
-
-    assert src_len == value_len
-
-    q = q.reshape(tgt_len, num_heads, head_dim)
-    k = k.reshape(src_len, num_heads, head_dim)
-    v = v.reshape(src_len, num_heads, head_dim)
-
-    if add_zero_attn:
-        zero_attn_shape = (1, num_heads, head_dim)
-        k_zeros = jnp.zeros(zero_attn_shape)
-        v_zeros = jnp.zeros(zero_attn_shape)
-
-        k = jnp.concatenate([k, k_zeros], axis=0)
-        v = jnp.concatenate([v, v_zeros], axis=0)
-
-        src_len += 1
-        value_len += 1
-
-    if bias_k is not None and bias_v is not None:
-        bias_k = bias_k.reshape(1, num_heads, head_dim)
-        bias_v = bias_v.reshape(1, num_heads, head_dim)
-
-        k = jnp.concatenate([k, bias_k], axis=0)
-        v = jnp.concatenate([v, bias_v], axis=0)
-
-        src_len += 1
-        value_len += 1
-
-    assert src_len == value_len
-
-    # [tgt_len, num_heads, head_dim] → [num_heads, tgt_len, head_dim]
-    q = jnp.transpose(q, (1, 0, 2))
-
-    # [src_len, num_heads, head_dim] → [num_heads, src_len, head_dim]
-    k = jnp.transpose(k, (1, 0, 2))
-    v = jnp.transpose(v, (1, 0, 2))
-
-    scale = jnp.sqrt(head_dim)
-    attn_output_weights = jnp.matmul(q, jnp.transpose(k, (0, 2, 1))) / scale
-
-    if key_padding_mask is not None:
-        padding_mask = key_padding_mask.reshape(1, 1, src_len)
-        padding_mask = jnp.repeat(padding_mask, num_heads, axis=0)
-        padding_mask = jnp.repeat(padding_mask, tgt_len, axis=1)
-        attn_output_weights = jnp.where(
-            padding_mask, float("-inf"), attn_output_weights
-        )
-
-    if attn_mask is not None:
-        # [tgt_len, src_len] -> [num_heads, tgt_len, src_len]
-        mask = attn_mask.reshape(1, tgt_len, src_len)
-        mask = jnp.repeat(mask, num_heads, axis=0)
-        attn_output_weights = attn_output_weights + mask
-
-    if is_causal:
-        causal_mask = jnp.triu(jnp.ones((tgt_len, src_len)), k=1)
-        causal_mask = (causal_mask == 1).reshape(1, tgt_len, src_len)
-        causal_mask = jnp.repeat(causal_mask, num_heads, axis=0)
-        attn_output_weights = jnp.where(causal_mask, float("-inf"), attn_output_weights)
-
-    # [num_heads, tgt_len, src_len]
-    attn_output_weights = jax.nn.softmax(attn_output_weights, axis=-1)
-
-    if dropout_p > 0.0 and not inference:
-        assert dropout_key is not None, (
-            "dropout_key required because dropout_p > 0.0 and training"
-        )
-        dropout_mask = jax.random.bernoulli(
-            dropout_key, 1 - dropout_p, attn_output_weights.shape
-        )
-        scale = 1.0 / (1.0 - dropout_p)
-        attn_output_weights = attn_output_weights * dropout_mask * scale
-
-    attn_output = jnp.matmul(attn_output_weights, v)
-    attn_output = jnp.transpose(attn_output, (1, 0, 2))
-    attn_output = attn_output.reshape(tgt_len, d_model)
-
-    assert out_proj_weight is not None, "out_proj_weight must be provided"
-    attn_output = attn_output @ out_proj_weight.T
-
-    if out_proj_bias is not None:
-        attn_output = attn_output + out_proj_bias
-
-    if need_weights:
-        if average_attn_weights:
-            attn_output_weights = attn_output_weights.mean(axis=0)
-        return attn_output, attn_output_weights
-    else:
-        return attn_output, None
-
-
-def create_attn_mask(pad_H, pad_W, window_size, shift_size, dtype: Any | None = None):
-    if dtype is None:
-        dtype = default_floating_dtype()
-    assert dtype is not None
-    h_boundaries = jnp.array([pad_H - window_size[0], pad_H - shift_size[0]])
-    w_boundaries = jnp.array([pad_W - window_size[1], pad_W - shift_size[1]])
-
-    h_boundaries = jnp.sort(h_boundaries)
-    w_boundaries = jnp.sort(w_boundaries)
-
-    ii, jj = jnp.indices((pad_H, pad_W))  # ii for rows, jj for columns
-
-    row_region_idx = jnp.searchsorted(h_boundaries, ii, side="right")
-    col_region_idx = jnp.searchsorted(w_boundaries, jj, side="right")
-
-    num_col_regions = len(w_boundaries) + 1
-    attn_mask = row_region_idx * num_col_regions + col_region_idx
-
-    return attn_mask.astype(dtype)
-
-
-def shifted_window_attention(
-    x: Float[Array, "H W C"],
-    qkv_weight: Float[Array, "in_dim out_dim"],
-    proj_weight: Float[Array, "out_dim out_dim"],
-    relative_position_bias: Array,
-    window_size: list[int],
-    num_heads: int,
-    shift_size: list[int],
-    attention_dropout: float = 0.0,
-    dropout: float = 0.0,
-    qkv_bias: Array | None = None,
-    proj_bias: Array | None = None,
-    logit_scale: Array | None = None,
-    inference: bool = False,
-    key: PRNGKeyArray | None = None,
-) -> Float[Array, "H W C"]:
-    if not inference and key is None:
-        raise ValueError("Need key when in training mode")
-    H, W, C = x.shape
-    to_pad_W = (window_size[1] - W % window_size[1]) % window_size[1]
-    to_pad_H = (window_size[0] - H % window_size[0]) % window_size[0]
-    x = jnp.pad(x, ((0, to_pad_H), (0, to_pad_W), (0, 0)))
-    pad_H, pad_W, _ = x.shape
-
-    shift_size = shift_size.copy()
-    if window_size[0] >= pad_H:
-        shift_size[0] = 0
-    if window_size[1] >= pad_W:
-        shift_size[1] = 0
-
-    # cyclic shift
-    if sum(shift_size) > 0:
-        x = jnp.roll(x, shift=(-shift_size[0], -shift_size[1]), axis=(0, 1))
-
-    # partition windows
-    num_windows = (pad_H // window_size[0]) * (pad_W // window_size[1])
-    x = jnp.reshape(
-        x,
-        (
-            pad_H // window_size[0],
-            window_size[0],
-            pad_W // window_size[1],
-            window_size[1],
-            C,
-        ),
-    )
-    x = jnp.transpose(x, (0, 2, 1, 3, 4)).reshape(
-        num_windows, window_size[0] * window_size[1], C
-    )
-
-    # multi-head attention
-    if logit_scale is not None and qkv_bias is not None:
-        length = qkv_bias.size // 3
-        qkv_bias = qkv_bias.at[length : 2 * length].set(0.0)
-
-    def linear(x: Array, weight: Array, bias: Array | None):
-        output = x @ jnp.transpose(weight)  # (in,) @ (in, out) -> (out,)
-        if bias is not None:
-            output = output + bias
-        return output
-
-    linear_pt = functools.partial(linear, weight=qkv_weight, bias=qkv_bias)
-
-    qkv = eqx.filter_vmap(eqx.filter_vmap(linear_pt))(x)
-    win_size, patches, _ = qkv.shape
-    qkv = jnp.transpose(
-        qkv.reshape(win_size, patches, 3, num_heads, C // num_heads), (2, 0, 3, 1, 4)
-    )
-    q, k, v = qkv[0], qkv[1], qkv[2]
-    if logit_scale is not None:
-        # cosine attention
-        attn = normalize(q, axis=-1) @ jnp.transpose(
-            normalize(k, axis=-1), (0, 1, 3, 2)
-        )
-        # Clamp the logit scale exponent for stability
-        logit_scale = jnp.exp(jnp.minimum(logit_scale, jnp.log(jnp.array(100.0))))
-        attn = attn * logit_scale
-    else:
-        q = q * (C // num_heads) ** -0.5
-        # attn = q @ (jnp.transpose(normalize(k, axis=-1), (0, 1, 3, 2))) # Incorrect
-        attn = q @ jnp.transpose(k, (0, 1, 3, 2))  # Corrected: q @ k.T
-
-    # add relative position bias
-    attn = attn + relative_position_bias
-
-    if sum(shift_size) > 0:
-        attn_mask = create_attn_mask(pad_H, pad_W, window_size, shift_size)
-        attn_mask = attn_mask.reshape(
-            pad_H // window_size[0],
-            window_size[0],
-            pad_W // window_size[1],
-            window_size[1],
-        )
-        attn_mask = jnp.transpose(attn_mask, (0, 2, 1, 3)).reshape(
-            num_windows, window_size[0] * window_size[1]
-        )
-        attn_mask = jnp.expand_dims(attn_mask, axis=1) - jnp.expand_dims(
-            attn_mask, axis=2
-        )
-        attn_mask = jnp.where(attn_mask == 0, 0.0, -100.0)
-
-        attn = attn + attn_mask[:, None, :, :]
-
-    attn = jax.nn.softmax(attn, axis=-1)
-    if not inference:
-        assert key is not None, "key must be given if not inference"
-        key, subkey = jax.random.split(key)
-        attn = dropout_fn(attn, p=attention_dropout, inference=inference, key=subkey)
-
-    x = jnp.transpose(attn @ v, (0, 2, 1, 3)).reshape(
-        num_windows, window_size[0] * window_size[1], C
-    )
-    linear_pt_proj = functools.partial(linear, weight=proj_weight, bias=proj_bias)
-
-    x = eqx.filter_vmap(eqx.filter_vmap(linear_pt_proj))(x)
-    if not inference:
-        assert key is not None, "key must be given if not inference"
-        key, subkey = jax.random.split(key)
-        x = dropout_fn(x, p=dropout, inference=inference, key=subkey)
-
-    # reverse windows
-    x = x.reshape(
-        pad_H // window_size[0],
-        pad_W // window_size[1],
-        window_size[0],
-        window_size[1],
-        C,
-    )
-    x = jnp.transpose(x, (0, 2, 1, 3, 4)).reshape(pad_H, pad_W, C)
-
-    # reverse cyclic shift
-    if sum(shift_size) > 0:
-        x = jnp.roll(x, shift=(shift_size[0], shift_size[1]), axis=(0, 1))
-
-    # unpad features
-    x = x[:H, :W, :]
-    return x

jaxonlayers/functions/initialization.py

@@ -1,34 +0,0 @@
-import equinox as eqx
-import jax
-from jaxtyping import PRNGKeyArray, PyTree
-
-
-def kaiming_init_conv2d(model: PyTree, state: eqx.nn.State, key: PRNGKeyArray):
-    """Applies Kaiming He normal initialization to Conv2d weights."""
-    # Filter function to identify Conv2d layers
-    is_conv2d = lambda x: isinstance(x, eqx.nn.Conv2d)
-
-    # Function to get weights (leaves) based on the filter
-    def get_weights(model):
-        return [
-            x.weight for x in jax.tree.leaves(model, is_leaf=is_conv2d) if is_conv2d(x)
-        ]
-
-    # Get the list of current weights
-    weights = get_weights(model)
-    if not weights:  # If no Conv2d layers found
-        return model, state
-
-    # Create new weights using He initializer
-    initializer = jax.nn.initializers.he_normal()
-    # Split key for each weight tensor
-    subkeys = jax.random.split(key, len(weights))
-    new_weights = [
-        initializer(subkeys[i], w.shape, w.dtype)  # Use original weight's dtype
-        for i, w in enumerate(weights)
-    ]
-
-    # Replace old weights with new weights in the model pytree
-    model = eqx.tree_at(get_weights, model, new_weights)
-
-    return model, state

jaxonlayers/functions/masking.py

@@ -1,58 +0,0 @@
-import jax.numpy as jnp
-
-
-def canonical_mask(
-    mask,
-    mask_name,
-    other_name="",
-    other_type=None,
-    target_type=jnp.float32,
-    other_mask=None,
-    check_other=True,
-):
-    if mask is None:
-        return None
-    if mask.dtype == bool:
-        additive_mask = jnp.where(mask, -jnp.inf, 0.0).astype(target_type)
-        return additive_mask
-    elif jnp.issubdtype(mask.dtype, jnp.integer) or jnp.issubdtype(
-        mask.dtype, jnp.floating
-    ):
-        return mask.astype(target_type)
-    else:
-        raise TypeError(
-            f"{mask_name} must be bool, int, or float tensor, but got {mask.dtype}"
-        )
-
-
-def canonical_key_padding_mask(
-    key_padding_mask, attn_mask=None, query_dtype=jnp.float32
-):
-    """Wrapper for canonicalizing key_padding_mask"""
-    return canonical_mask(
-        mask=key_padding_mask,
-        mask_name="key_padding_mask",
-        other_name="attn_mask",
-        other_mask=attn_mask,
-        target_type=query_dtype,
-    )
-
-
-def canonical_attn_mask(attn_mask, query_dtype=jnp.float32):
-    """Wrapper for canonicalizing attn_mask"""
-    return canonical_mask(
-        mask=attn_mask,
-        mask_name="attn_mask",
-        other_type=None,
-        other_name="",
-        target_type=query_dtype,
-        check_other=False,
-    )
-
-
-def build_attention_mask(context_length: int):
-    mask = jnp.tril(jnp.zeros((context_length, context_length)))
-    upper = jnp.triu(jnp.full((context_length, context_length), float("-inf")), k=1)
-
-    mask = mask + upper
-    return mask

jaxonlayers/functions/normalization.py

@@ -1,9 +0,0 @@
-import jax.numpy as jnp
-
-
-def normalize(x, p=2, axis=1, eps=1e-12):
-    norm = jnp.linalg.norm(x, ord=p, axis=axis, keepdims=True)
-    norm = jnp.maximum(norm, eps)
-    output = x / norm
-
-    return output

jaxonlayers/functions/regularization.py

@@ -1,47 +0,0 @@
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array, PRNGKeyArray
-
-
-def stochastic_depth(
-    input: Array,
-    p: float,
-    mode: str,
-    inference: bool,
-    key: PRNGKeyArray,
-) -> Array:
-    if p < 0.0 or p > 1.0:
-        raise ValueError(f"drop probability has to be between 0 and 1, but got {p}")
-    if mode not in ["batch", "row"]:
-        raise ValueError(f"mode has to be either 'batch' or 'row', but got {mode}")
-    if inference or p == 0.0:
-        return input
-    survival_rate = 1.0 - p
-    if mode == "row":
-        size = [input.shape[0]] + [1] * (input.ndim - 1)
-    else:
-        size = [1] * input.ndim
-    noise = jax.random.bernoulli(key, p=survival_rate, shape=size).astype(input.dtype)
-    if survival_rate > 0.0:
-        noise = noise / survival_rate
-    return input * noise
-
-
-def dropout(
-    x: Array,
-    p: float,
-    inference: bool,
-    key: PRNGKeyArray | None = None,
-) -> Array:
-    if isinstance(p, (int, float)) and p == 0:
-        inference = True
-    if inference:
-        return x
-    elif key is None:
-        raise RuntimeError(
-            "Dropout requires a key when running in non-deterministic mode."
-        )
-    else:
-        q = 1 - jax.lax.stop_gradient(p)
-        mask = jax.random.bernoulli(key, q, x.shape)
-        return jnp.where(mask, x / q, 0)

jaxonlayers/functions/state_space.py

@@ -1,30 +0,0 @@
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array, Float
-
-
-def selective_scan(
-    x: Float[Array, "seq_length d_inner"],
-    delta: Float[Array, "seq_length d_inner"],
-    A: Float[Array, "d_inner d_state"],
-    B: Float[Array, "seq_length d_state"],
-    C: Float[Array, "seq_length d_state"],
-    D: Float[Array, " d_inner"],
-) -> Float[Array, "seq_length d_inner"]:
-    L, d_inner = x.shape
-    _, d_state = A.shape
-    delta_A = jnp.exp(jnp.einsum("l d,d n -> l d n", delta, A))
-    delta_B_u = jnp.einsum("l d,l n,l d -> l d n", delta, B, x)
-
-    x_res = jnp.zeros(shape=(d_inner, d_state))
-
-    def step(x, i):
-        x = delta_A[i] * x + delta_B_u[i]
-
-        y = jnp.einsum("d n,n -> d", x, C[i, :])
-        return x, y
-
-    _, ys = jax.lax.scan(step, x_res, jnp.arange(L))
-
-    ys = ys + x * D
-    return ys

jaxonlayers/functions/utils.py

@@ -1,9 +0,0 @@
-import jax
-import jax.numpy as jnp
-
-
-def default_floating_dtype():
-    if jax.config.jax_enable_x64:  # pyright: ignore
-        return jnp.float64
-    else:
-        return jnp.float32

jaxonlayers/layers/__init__.py

@@ -1,16 +0,0 @@
-from .attention import MultiheadAttention, SqueezeExcitation
-from .convolution import ConvNormActivation
-from .normalization import BatchNorm, LayerNorm, LocalResponseNormalization
-from .regularization import StochasticDepth
-from .state_space import SelectiveStateSpace
-
-__all__ = [
-    "BatchNorm",
-    "LocalResponseNormalization",
-    "MultiheadAttention",
-    "SelectiveStateSpace",
-    "SqueezeExcitation",
-    "StochasticDepth",
-    "ConvNormActivation",
-    "LayerNorm",
-]

jaxonlayers/layers/abstract.py

@@ -1,16 +0,0 @@
-import abc
-
-import equinox as eqx
-from jaxtyping import Array
-
-
-class AbstractNorm(eqx.Module):
-    @abc.abstractmethod
-    def __call__(self, x: Array, *_, **__) -> Array: ...
-
-
-class AbstractNormStateful(eqx.nn.StatefulLayer):
-    @abc.abstractmethod
-    def __call__(
-        self, x: Array, state: eqx.nn.State, *_, **__
-    ) -> tuple[Array, eqx.nn.State]: ...