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

jaxonlayers/functions/__init__.py

@@ -0,0 +1,27 @@
+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

@@ -0,0 +1,374 @@
+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/embedding.py

@@ -0,0 +1,21 @@
+import jax.numpy as jnp
+from jaxtyping import Array, Float, Int
+
+
+def sinusoidal_embedding(
+    t: Int[Array, ""], embedding_size: int
+) -> Float[Array, " embedding_size"]:
+    if embedding_size % 2 != 0:
+        raise ValueError(f"Embedding size must be even, but got {embedding_size}")
+
+    half_dim = embedding_size // 2
+    embedding_freqs = jnp.exp(
+        -jnp.log(10000)
+        * jnp.arange(start=0, stop=half_dim, dtype=jnp.float32)
+        / half_dim
+    )
+
+    time_args = t * embedding_freqs
+    embedding = jnp.concatenate([jnp.sin(time_args), jnp.cos(time_args)])
+
+    return embedding

jaxonlayers/functions/initialization.py

@@ -0,0 +1,34 @@
+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

@@ -0,0 +1,58 @@
+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

@@ -0,0 +1,9 @@
+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

@@ -0,0 +1,47 @@
+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

@@ -0,0 +1,79 @@
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+
+def selective_scan(
+    u: 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_inner d_state"],
+    C: Float[Array, "seq_length d_inner d_state"],
+    D: Float[Array, " d_inner"],
+    chunk_size: int = 128,
+) -> Float[Array, "seq_length d_inner"]:
+    deltaA = jnp.exp(jnp.einsum("l d, d n -> l d n", delta, A))
+    deltaB_u = jnp.einsum("l d, l d n, l d -> l d n", delta, B, u)
+
+    seq_len, d_inner = u.shape
+    d_state = A.shape[1]
+
+    num_chunks = (seq_len + chunk_size - 1) // chunk_size
+    padded_len = num_chunks * chunk_size
+
+    pad_len = padded_len - seq_len
+    deltaA_padded = jnp.pad(deltaA, ((0, pad_len), (0, 0), (0, 0)))
+    deltaB_u_padded = jnp.pad(deltaB_u, ((0, pad_len), (0, 0), (0, 0)))
+    C_padded = jnp.pad(C, ((0, pad_len), (0, 0), (0, 0)))
+
+    deltaA_chunked = deltaA_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
+    deltaB_u_chunked = deltaB_u_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
+    C_chunked = C_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
+
+    def intra_chunk_step(h_prev, scan_inputs):
+        deltaA_i, deltaB_u_i, C_i = scan_inputs
+        h_i = deltaA_i * h_prev + deltaB_u_i
+        y_i = jnp.einsum("d n, d n -> d", h_i, C_i)
+        return h_i, y_i
+
+    h0 = jnp.zeros((d_inner, d_state))
+
+    _, y_chunks = jax.vmap(jax.lax.scan, in_axes=(None, None, 0))(
+        intra_chunk_step, h0, (deltaA_chunked, deltaB_u_chunked, C_chunked)
+    )
+
+    def inter_chunk_step(carry_prev, scan_inputs):
+        A_prev, h_prev = carry_prev
+        deltaA_i, deltaB_u_i = scan_inputs
+
+        A_new = deltaA_i * A_prev
+        h_new = deltaA_i * h_prev + deltaB_u_i
+
+        return (A_new, h_new), (A_new, h_new)
+
+    A_carry_initial = jnp.ones((d_inner, d_state))
+    h_carry_initial = jnp.zeros((d_inner, d_state))
+    initial_carry = (A_carry_initial, h_carry_initial)
+
+    scan_inputs = (deltaA_chunked[:, -1], deltaB_u_chunked[:, -1])
+
+    _, (A_carry, h_carry) = jax.lax.scan(inter_chunk_step, initial_carry, scan_inputs)
+
+    A_carry = jnp.roll(A_carry, 1, axis=0)
+    h_carry = jnp.roll(h_carry, 1, axis=0)
+    A_carry = A_carry.at[0].set(jnp.ones_like(A_carry[0]))
+    h_carry = h_carry.at[0].set(jnp.zeros_like(h_carry[0]))
+
+    h_carry_broadcast = jnp.expand_dims(h_carry, axis=1)
+    h_correction = deltaA_chunked * h_carry_broadcast
+    y_carry = jnp.einsum("csdn, csdn -> csd", C_chunked, h_correction)
+
+    y_final = y_chunks + y_carry
+
+    y_final = y_final.reshape(padded_len, d_inner)
+
+    y_unpadded = y_final[:seq_len]
+
+    output = y_unpadded.real + u * D
+
+    return output