xax 0.1.11__tar.gz → 0.1.12__tar.gz

{xax-0.1.11/xax.egg-info → xax-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.11
+Version: 0.1.12
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.1.11 → xax-0.1.12}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.11"
+__version__ = "0.1.12"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -43,15 +43,14 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
+    "rotate_vector_by_quat",
     "cross_entropy",
     "cast_norm_type",
     "get_norm",
     "is_master",
+    "BaseSSMBlock",
     "DiagSSMBlock",
-    "DiscreteTimeS4",
-    "S4",
-    "S4Layer",
-    "S6Layer",
+    "SSM",
     "SSMBlock",
     "BaseLauncher",
     "CliLauncher",
@@ -203,15 +202,14 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
+    "rotate_vector_by_quat": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.norm",
     "get_norm": "nn.norm",
     "is_master": "nn.parallel",
+    "BaseSSMBlock": "nn.ssm",
     "DiagSSMBlock": "nn.ssm",
-    "DiscreteTimeS4": "nn.ssm",
-    "S4": "nn.ssm",
-    "S4Layer": "nn.ssm",
-    "S6Layer": "nn.ssm",
+    "SSM": "nn.ssm",
     "SSMBlock": "nn.ssm",
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
@@ -364,11 +362,12 @@ if IMPORT_ALL or TYPE_CHECKING:
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         quat_to_euler,
+        rotate_vector_by_quat,
     )
     from xax.nn.losses import cross_entropy
     from xax.nn.norm import NormType, cast_norm_type, get_norm
     from xax.nn.parallel import is_master
-    from xax.nn.ssm import S4, DiagSSMBlock, DiscreteTimeS4, S4Layer, S6Layer, SSMBlock
+    from xax.nn.ssm import SSM, BaseSSMBlock, DiagSSMBlock, SSMBlock
     from xax.task.base import RawConfigType
     from xax.task.launchers.base import BaseLauncher
     from xax.task.launchers.cli import CliLauncher

{xax-0.1.11 → xax-0.1.12}/xax/nn/geom.py

@@ -99,3 +99,60 @@ def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -
 
     # Note: We're rotating [0,0,-1], so we negate gz to match the expected direction
     return jnp.concatenate([gx, gy, -gz], axis=-1)
+
+
+def rotate_vector_by_quat(vector: jax.Array, quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+    """Rotates a vector by a quaternion.
+
+    Args:
+        vector: The vector to rotate, shape (*, 3).
+        quat: The quaternion to rotate by, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    Returns:
+        The rotated vector, shape (*, 3).
+    """
+    # Normalize quaternion
+    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
+    w, x, y, z = jnp.split(quat, 4, axis=-1)
+
+    # Extract vector components
+    vx, vy, vz = jnp.split(vector, 3, axis=-1)
+
+    # Terms for x component
+    xx = (
+        w * w * vx
+        + 2 * y * w * vz
+        - 2 * z * w * vy
+        + x * x * vx
+        + 2 * y * x * vy
+        + 2 * z * x * vz
+        - z * z * vx
+        - y * y * vx
+    )
+
+    # Terms for y component
+    yy = (
+        2 * x * y * vx
+        + y * y * vy
+        + 2 * z * y * vz
+        + 2 * w * z * vx
+        - z * z * vy
+        + w * w * vy
+        - 2 * w * x * vz
+        - x * x * vy
+    )
+
+    # Terms for z component
+    zz = (
+        2 * x * z * vx
+        + 2 * y * z * vy
+        + z * z * vz
+        - 2 * w * y * vx
+        + w * w * vz
+        + 2 * w * x * vy
+        - y * y * vz
+        - x * x * vz
+    )
+
+    return jnp.concatenate([xx, yy, zz], axis=-1)

xax-0.1.12/xax/nn/ssm.py

@@ -0,0 +1,316 @@
+"""State space models."""
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, PRNGKeyArray
+
+
+def glorot(key: PRNGKeyArray, shape: tuple[int, ...]) -> Array:
+    return jax.random.uniform(key, shape, minval=-1.0, maxval=1.0) * jnp.sqrt(2 / sum(shape))
+
+
+class BaseSSMBlock(eqx.Module, ABC):
+    @abstractmethod
+    def forward(self, h: Array, x: Array) -> Array: ...
+
+    @abstractmethod
+    def forward_sequence(self, x_seq: Array) -> Array: ...
+
+    @abstractmethod
+    def get_a_mat(self, x: Array) -> Array: ...
+
+    @abstractmethod
+    def get_b_mat(self, x: Array) -> Array: ...
+
+
+class SSMBlock(BaseSSMBlock):
+    a_mat: Array
+    b_mat: Array
+
+    def __init__(self, hidden_size: int, *, key: PRNGKeyArray) -> None:
+        key_a, key_b = jax.random.split(key)
+        self.a_mat = glorot(key_a, (hidden_size, hidden_size))
+        self.b_mat = glorot(key_b, (hidden_size, hidden_size))
+
+    def get_a_mat(self, x: Array) -> Array:
+        return self.a_mat
+
+    def get_b_mat(self, x: Array) -> Array:
+        return self.b_mat
+
+    def forward(self, h: Array, x: Array) -> Array:
+        """Perform a forward pass.
+
+        Args:
+            h: Hidden state of shape (H,).
+            x: Input of shape (H,).
+
+        Returns:
+            Hidden state of shape (H,).
+        """
+        a_mat = self.get_a_mat(x)
+        b_mat = self.get_b_mat(x)
+        h = a_mat @ h + b_mat.T @ x
+        return h
+
+    def forward_sequence(self, x_seq: Array) -> Array:
+        """Perform a forward pass across time.
+
+        Args:
+            x_seq: Input sequence of shape (T, H).
+
+        Returns:
+            Hidden state sequence of shape (T, H).
+        """
+
+        def step(h: Array, x: Array) -> tuple[Array, Array]:
+            h = self.forward(h, x)
+            return h, h
+
+        a_mat = self.get_a_mat(x_seq)
+        h_0 = jnp.zeros(a_mat.shape[0])
+        _, h_seq = jax.lax.scan(step, h_0, x_seq)
+        return h_seq
+
+
+class DiagSSMBlock(BaseSSMBlock):
+    a_diag: Array
+    b_mat: Array
+
+    def __init__(self, hidden_size: int, *, key: PRNGKeyArray) -> None:
+        keys = jax.random.split(key, 2)
+        self.a_diag = glorot(keys[0], (hidden_size,))
+        self.b_mat = glorot(keys[1], (hidden_size, hidden_size))
+
+    def get_a_mat(self, x: Array) -> Array:
+        return self.a_diag
+
+    def get_b_mat(self, x: Array) -> Array:
+        return self.b_mat
+
+    def forward(self, h: Array, x: Array) -> Array:
+        """Perform a forward pass.
+
+        Args:
+            h: Hidden state of shape (H,).
+            x: Input of shape (H,).
+
+        Returns:
+            Hidden state of shape (H,).
+        """
+        a_diag = self.get_a_mat(x)
+        b_mat = self.get_b_mat(x)
+        h = a_diag * h + b_mat.T @ x
+        return h
+
+    def forward_sequence(self, x_seq: Array, *, use_conv: bool = True, recursive_kernel_calc: bool = False) -> Array:
+        """Perform a potentially parallelized forward pass across time.
+
+        Args:
+            x_seq: Input sequence of shape (T, H).
+            use_conv: Whether to use convolution to compute the sequence.
+            recursive_kernel_calc: Whether to use a recursive kernel calculation.
+
+        Returns:
+            Hidden state sequence of shape (T, H).
+        """
+        if use_conv:
+            return self._forward_sequence_conv(x_seq, recursive_kernel_calc=recursive_kernel_calc)
+        else:
+            return self._forward_sequence_scan(x_seq)
+
+    def _get_kernel(self, x_seq: Array, length: int) -> Array:
+        """Returns the kernel with time as the final dimension."""
+        exponents = jnp.arange(length)
+        a_diag = self.get_a_mat(x_seq)
+        kernel = jnp.power(a_diag[:, None], exponents)  # (H, T)
+        kernel = kernel[:, None, :]  # (H, 1, T)
+        return kernel
+
+    def _get_kernel_recursive(self, x_seq: Array, length: int) -> Array:
+        """Returns the kernel with time as the final dimension."""
+        assert length % 2 == 0, "Length must be even."
+        a_diag = self.get_a_mat(x_seq)
+
+        def helper(length: int) -> tuple[Array, Array]:
+            """Returns the kernel and the sqrt of the diagonal."""
+            if length == 1:
+                return jnp.ones_like(a_diag)[:, None], a_diag[:, None]
+
+            half_length = length // 2
+            kernel_half, a_half = helper(half_length)
+            kernel = jnp.concatenate([kernel_half, a_half * kernel_half], axis=-1)
+            return kernel, a_half * a_half
+
+        kernel, a_diag = helper(length)
+        return kernel[:, None, :]  # (H, 1, L)
+
+    def _forward_sequence_conv(self, x_seq: Array, *, recursive_kernel_calc: bool = False) -> Array:
+        """Convolves x (T, H) across time using the kernel."""
+        seq_len, hidden_size = x_seq.shape
+        b_mat = self.get_b_mat(x_seq)
+
+        s = b_mat.T @ x_seq.T  # (H, T)
+        s_padded = jnp.pad(s, ((0, 0), (seq_len - 1, 0)))[None, :, :]  # (1, H, 2T-1)
+
+        if recursive_kernel_calc:
+            kernel = self._get_kernel_recursive(x_seq, seq_len)
+        else:
+            kernel = self._get_kernel(x_seq, seq_len)
+
+        kernel_flipped = jnp.flip(kernel, axis=-1)  # (H, 1, L)
+
+        conv_out = jax.lax.conv_general_dilated(
+            s_padded,
+            kernel_flipped,
+            window_strides=(1,),
+            padding="VALID",
+            dimension_numbers=("NCT", "OIT", "NCT"),  # convolving over time
+            feature_group_count=hidden_size,
+        )
+        conv_out = conv_out[0].T  # (T, H)
+        return conv_out
+
+    def _forward_sequence_scan(self, x_seq: Array) -> Array:
+        """Naively forward across time."""
+
+        def step(h: Array, x: Array) -> tuple[Array, Array]:
+            h = self.forward(h, x)
+            return h, h
+
+        a_diag = self.get_a_mat(x_seq)
+        h_0 = jnp.zeros(a_diag.shape[0])
+        _, h_seq = jax.lax.scan(step, h_0, x_seq)
+        return h_seq
+
+
+class DiscreteDiagSSMBlock(DiagSSMBlock):
+    delta: Array
+
+    def __init__(
+        self,
+        hidden_size: int,
+        *,
+        key: PRNGKeyArray,
+        init_delta: float = 1.0,
+        init_scale: float = 10.0,
+    ) -> None:
+        super().__init__(hidden_size, key=key)
+        self.delta = jnp.array(init_delta)
+
+        # A positive scale helps reduce the gradient at the start.
+        self.a_diag = jax.random.uniform(key, (hidden_size,), minval=-1.0, maxval=0.0) * init_scale
+
+    def get_a_mat(self, x: Array) -> Array:
+        """Discretize the diagonal matrix using zero-order hold."""
+        a_diag_discrete = jnp.exp(self.a_diag * self.delta)
+        return a_diag_discrete
+
+    def get_b_mat(self, x: Array) -> Array:
+        """Discretize the input matrix using zero-order hold."""
+        delta_a_diag = self.a_diag * self.delta
+        exp_a_diag = jnp.exp(delta_a_diag)
+        delta_a_inv = 1 / delta_a_diag
+        delta_b_mat = self.delta * self.b_mat
+
+        b_discrete = delta_a_inv * (exp_a_diag - 1) * delta_b_mat
+        return b_discrete
+
+
+class SSM(eqx.Module):
+    vocab_embedding: eqx.nn.Embedding
+    output_layer: eqx.nn.Linear
+    blocks: list[BaseSSMBlock]
+    num_layers: int = eqx.static_field()
+    hidden_size: int = eqx.static_field()
+    skip_connections: bool = eqx.static_field()
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        output_size: int,
+        num_layers: int,
+        block_type: Literal["diagonal", "full_rank"] = "full_rank",
+        skip_connections: bool = False,
+        discretize: bool = False,
+        *,
+        key: PRNGKeyArray,
+    ) -> None:
+        vocab_key, s4_key = jax.random.split(key, 2)
+        self.vocab_embedding = eqx.nn.Embedding(input_size, hidden_size, key=vocab_key)
+        self.output_layer = eqx.nn.Linear(hidden_size, output_size, key=key)
+
+        block_keys = jax.random.split(s4_key, num_layers)
+
+        def get_block(key: PRNGKeyArray) -> BaseSSMBlock:
+            match block_type:
+                case "diagonal":
+                    return (
+                        DiscreteDiagSSMBlock(hidden_size, key=key, init_delta=0.1)
+                        if discretize
+                        else DiagSSMBlock(hidden_size, key=key)
+                    )
+                case "full_rank":
+                    if discretize:
+                        raise ValueError("Full rank blocks do not support discretization due to instability.")
+                    return SSMBlock(hidden_size, key=key)
+                case _:
+                    raise ValueError(f"Unknown block type: {block_type}")
+
+        self.blocks = [get_block(block_keys[i]) for i in range(num_layers)]
+        self.skip_connections = skip_connections
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+
+    def __call__(self, hs: list[Array], x: Array) -> tuple[list[Array], Array]:
+        new_hs = []
+        for i, block in enumerate(self.blocks):
+            h = block.forward(hs[i], x)
+            new_hs.append(h)
+            xh = jax.nn.gelu(h)
+            x = xh + x if self.skip_connections else xh
+        y = self.output_layer(x)
+        return new_hs, y
+
+    def _embed_input(self, x: Array) -> Array:
+        """U is the input to the S4 cell."""
+        return self.vocab_embedding(x)
+
+    def predict_sequence(self, x_seq: Array) -> Array:
+        x_emb = jax.vmap(self._embed_input)(x_seq)
+        for block in self.blocks:
+            h = block.forward_sequence(x_emb)
+            # h = block.naive_forward_sequence(x_emb)
+            h = jax.nn.gelu(h)
+            x_emb = h + x_emb if self.skip_connections else h
+        y = jax.vmap(self.output_layer)(x_emb)
+        return y
+
+    def generate_sequence(self, prompt_seq: Array, max_len: int) -> Array:
+        hs = [jnp.zeros(self.hidden_size) for _ in range(self.num_layers)]
+        prompt_seq_embedded = jax.vmap(self._embed_input)(prompt_seq)
+
+        def encode_step(hs: list[Array], x: Array) -> tuple[list[Array], Array]:
+            hs, y = self(hs, x)
+            return hs, y
+
+        def decode_step(
+            carry: tuple[list[Array], Array, PRNGKeyArray],
+            _: None,
+        ) -> tuple[tuple[list[Array], Array, PRNGKeyArray], Array]:
+            hs, last_token, rng = carry
+            token_embedded = self._embed_input(last_token)
+            hs, y = self(hs, token_embedded)
+            token = jax.random.categorical(rng, y)
+            rng = jax.random.split(rng)[0]
+            return (hs, token, rng), token
+
+        hs, _ = jax.lax.scan(encode_step, hs, prompt_seq_embedded)
+        _, sequence = jax.lax.scan(decode_step, (hs, prompt_seq[-1], jax.random.PRNGKey(0)), None, length=max_len)
+
+        return sequence

{xax-0.1.11 → xax-0.1.12}/xax/task/mixins/train.py

@@ -218,26 +218,32 @@ class TrainMixin(
         state = super().on_step_end(state)
         return state.replace(elapsed_time_s=time.time() - state.start_time_s)
 
-    def log_train_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
+    def log_train_step(
+        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+    ) -> None:
         """Override this function to do logging during the training phase.
 
         This function is called after the model forward pass and before the
         backward pass. It is called in the training phase.
 
         Args:
+            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
             metrics: The metrics for the current batch.
             state: The current training state.
         """
 
-    def log_valid_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
+    def log_valid_step(
+        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+    ) -> None:
         """Override this function to do logging during the validation phase.
 
         This function is called after the model forward pass. It is called in
         the validation phase.
 
         Args:
+            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
             metrics: The metrics for the current batch.
@@ -251,7 +257,9 @@ class TrainMixin(
             for k, v in d.items():
                 self.logger.log_scalar(k, v, namespace=ns)
 
-    def log_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
+    def log_step(
+        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+    ) -> None:
         phase = state.phase
 
         for k, v in metrics.items():
@@ -265,9 +273,9 @@ class TrainMixin(
         # Delegate to the appropriate logging function based on the phase.
         match phase:
             case "train":
-                self.log_train_step(batch, output, metrics, state)
+                self.log_train_step(model, batch, output, metrics, state)
             case "valid":
-                self.log_valid_step(batch, output, metrics, state)
+                self.log_valid_step(model, batch, output, metrics, state)
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
@@ -579,7 +587,7 @@ class TrainMixin(
                 )
 
                 output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
-                self.log_step(valid_batch, output, metrics, state)
+                self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
 
             state = self.on_step_start(state)
             train_batch = next(train_pf)
@@ -597,7 +605,7 @@ class TrainMixin(
                 batch=train_batch,
                 state=state,
             )
-            self.log_step(train_batch, output, metrics, state)
+            self.log_step(eqx.combine(model_arr, model_static), train_batch, output, metrics, state)
 
             state = self.on_step_end(state)
 

{xax-0.1.11 → xax-0.1.12/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.11
+Version: 0.1.12
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

xax-0.1.11/xax/nn/ssm.py

@@ -1,296 +0,0 @@
-"""State space models."""
-
-from abc import ABC, abstractmethod
-from typing import Literal
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array, PRNGKeyArray
-
-
-def glorot(key: PRNGKeyArray, shape: tuple[int, ...]) -> Array:
-    return jax.random.uniform(key, shape, minval=-1.0, maxval=1.0) * jnp.sqrt(2 / sum(shape))
-
-
-class DiscreteTimeS4(eqx.Module):
-    a: Array
-    B: Array
-    C: Array
-    proj_in: eqx.nn.Linear
-    proj_out: eqx.nn.Linear
-
-    def __init__(
-        self,
-        hidden_size: int,
-        projection_size: int,
-        input_size: int,
-        output_size: int,
-        *,
-        key: PRNGKeyArray,
-    ) -> None:
-        self.a = jax.nn.initializers.glorot_uniform()(key, (hidden_size,))
-        self.B = jax.nn.initializers.glorot_uniform()(key, (projection_size, hidden_size))
-        self.C = jax.nn.initializers.glorot_uniform()(key, (hidden_size, projection_size))
-        self.proj_in = eqx.nn.Linear(input_size, projection_size, key=key)
-        self.proj_out = eqx.nn.Linear(projection_size, output_size, key=key)
-
-    def __call__(self, h: Array, x: Array) -> tuple[Array, Array]:
-        h = self.a * h + self.B.T @ x
-        y = self.C.T @ h
-        return h, y
-
-    def predict_sequence(self, x_seq: Array) -> Array:
-        x_proj = jax.vmap(lambda x: jax.nn.relu(self.proj_in(x)))(x_seq)
-        h = jnp.zeros(self.a.shape[0])
-
-        def scan_fn(h: Array, x: Array) -> tuple[Array, Array]:
-            h = self.a * h + self.B.T @ x
-            y = self.C.T @ h
-            return h, y
-
-        _, y_seq = jax.lax.scan(scan_fn, h, x_proj)
-        y_out = jax.vmap(self.proj_out)(y_seq)
-        return y_out
-
-
-class S4Layer(eqx.Module):
-    a: Array
-    B: Array
-    C: Array
-    proj_in: eqx.nn.Linear
-    proj_out: eqx.nn.Linear
-    delta: Array
-
-    def __init__(
-        self,
-        hidden_size: int,
-        projection_size: int,
-        input_size: int,
-        output_size: int,
-        *,
-        key: PRNGKeyArray,
-    ) -> None:
-        self.a = jax.nn.initializers.glorot_uniform()(key, (hidden_size,))
-        self.B = jax.nn.initializers.glorot_uniform()(key, (projection_size, hidden_size))
-        self.C = jax.nn.initializers.glorot_uniform()(key, (hidden_size, projection_size))
-        self.proj_in = eqx.nn.Linear(input_size, projection_size, key=key)
-        self.proj_out = eqx.nn.Linear(projection_size, output_size, key=key)
-        self.delta = jax.random.uniform(key, (hidden_size,))
-
-    def __call__(self, h: Array, x: Array) -> tuple[Array, Array]:
-        delta_a = self.delta * self.a
-        a_bar = jnp.exp(delta_a)
-        b_bar = jnp.linalg.inv(delta_a) * (a_bar - 1) @ (self.delta * self.B)
-        h = a_bar * h + b_bar.T @ x
-        y = self.C.T @ h
-        return h, y
-
-    def predict_sequence(self, x_seq: Array) -> Array:
-        x_proj = jax.vmap(lambda x: jax.nn.gelu(self.proj_in(x)))(x_seq)
-        h = jnp.zeros(self.a.shape[0])
-
-        def scan_fn(h: Array, x: Array) -> tuple[Array, Array]:
-            h = self.a * h + self.B.T @ x
-            y = self.C.T @ h
-            return h, y
-
-        _, y_seq = jax.lax.scan(scan_fn, h, x_proj)
-        y_out = jax.vmap(self.proj_out)(y_seq)
-        return y_out
-
-
-class S6Layer(eqx.Module):
-    a: Array
-    B: Array
-    C: Array
-    proj_in: eqx.nn.Linear
-    proj_out: eqx.nn.Linear
-    delta: Array
-
-    def __init__(
-        self,
-        hidden_size: int,
-        projection_size: int,
-        input_size: int,
-        output_size: int,
-        *,
-        key: PRNGKeyArray,
-    ) -> None:
-        self.a = jax.nn.initializers.glorot_uniform()(key, (hidden_size,))
-        self.B = jax.nn.initializers.glorot_uniform()(key, (projection_size, hidden_size))
-        self.C = jax.nn.initializers.glorot_uniform()(key, (hidden_size, projection_size))
-        self.proj_in = eqx.nn.Linear(input_size, projection_size, key=key)
-        self.proj_out = eqx.nn.Linear(projection_size, output_size, key=key)
-        self.delta = jax.random.uniform(key, (hidden_size,))
-
-    def __call__(self, h: Array, x: Array) -> tuple[Array, Array]:
-        h = self.a * h + self.B.T @ x
-        y = self.C.T @ h
-        return h, y
-
-    def predict_sequence(self, x_seq: Array) -> Array:
-        x_proj = jax.vmap(lambda x: jax.nn.gelu(self.proj_in(x)))(x_seq)
-        h = jnp.zeros(self.a.shape[0])
-
-        def scan_fn(h: Array, x: Array) -> tuple[Array, Array]:
-            h = self.a * h + self.B.T @ x
-            y = self.C.T @ h
-            return h, y
-
-        _, y_seq = jax.lax.scan(scan_fn, h, x_proj)
-        y_out = jax.vmap(self.proj_out)(y_seq)
-        return y_out
-
-
-class BaseSSMBlock(eqx.Module, ABC):
-    @abstractmethod
-    def forward(self, h: Array, x: Array) -> Array:
-        pass
-
-
-class SSMBlock(BaseSSMBlock):
-    a_mat: Array
-    b_mat: Array
-
-    def __init__(self, hidden_size: int, *, key: PRNGKeyArray) -> None:
-        key_a, key_b = jax.random.split(key)
-        self.a_mat = glorot(key_a, (hidden_size, hidden_size))
-        self.b_mat = glorot(key_b, (hidden_size, hidden_size))
-
-    def forward(self, h: Array, x: Array) -> Array:
-        h = self.a_mat @ h + self.b_mat.T @ x
-        return h
-
-    def get_kernel(self, length: int) -> Array:
-        return self.a_mat
-
-
-class DiagSSMBlock(BaseSSMBlock):
-    a_mat: Array
-    b_mat: Array
-
-    def __init__(self, hidden_size: int, *, key: PRNGKeyArray) -> None:
-        keys = jax.random.split(key, 2)
-        self.a_mat = glorot(keys[0], (hidden_size,))
-        self.b_mat = glorot(keys[1], (hidden_size, hidden_size))
-
-    def forward(self, h: Array, x: Array) -> Array:
-        h = self.a_mat * h + self.b_mat.T @ x
-        h = jax.nn.tanh(h)
-        return h
-
-    def get_kernel(self, length: int) -> Array:
-        """Returns the kernel with time as the final dimension."""
-        exponents = jnp.arange(length)
-        kernel = jnp.power(self.a_mat[:, None], exponents)  # (H, L)
-        kernel = kernel[:, None, :]  # (H, 1, L)
-        return kernel
-
-    def forward_across_time(self, x: Array) -> Array:
-        """Convolves x (T, H) across time using the kernel."""
-        tsz, nhid = x.shape
-
-        # Compute s = x @ U.T + b, with shape (N, T, H)
-        s = self.b_mat.T @ x
-        s = s.T  # (H, T)
-
-        kernel = self.get_kernel(tsz)  # (H, 1, T)
-        kernel_flipped = jnp.flip(kernel, axis=-1)
-
-        # Pad s on the left along the time axis (pad length T-1)
-        s_padded = jnp.pad(s, ((0, 0), (0, 0), (tsz - 1, 0)))
-
-        # Perform depthwise (grouped) 1D convolution.
-        # We use input shape (N, H, L) and kernel shape (H, 1, T) with feature_group_count=H.
-        # The dimension_numbers are chosen so that the channel dimension is second.
-        conv_out = jax.lax.conv_general_dilated(
-            s_padded,
-            kernel_flipped,
-            window_strides=(1,),
-            padding="VALID",
-            dimension_numbers=("NCH", "OIH", "NCH"),
-            feature_group_count=nhid,
-        )
-        # conv_out has shape (N, H, T); transpose to (N, T, H)
-        conv_out = jnp.transpose(conv_out, (0, 2, 1))
-        return conv_out
-
-    def naive_forward_accross_time(self, x: Array) -> Array:
-        """Naively forward across time."""
-
-        def step(h: Array, x: Array) -> tuple[Array, Array]:
-            h = self.forward(h, x)
-            return h, h
-
-        h_0 = jnp.zeros(self.a_mat.shape[0])
-        _, h_seq = jax.lax.scan(step, h_0, x)
-        return h_seq
-
-
-class S4(eqx.Module):
-    vocab_embedding: eqx.nn.Embedding
-    proj_in: eqx.nn.Linear
-    proj_out: eqx.nn.Linear
-    blocks: list[BaseSSMBlock]
-    num_layers: int = eqx.static_field()
-    hidden_size: int = eqx.static_field()
-    skip_connections: bool = eqx.static_field()
-
-    def __init__(
-        self,
-        input_size: int,
-        hidden_size: int,
-        output_size: int,
-        num_layers: int,
-        block_type: Literal["ssm", "diag"] = "ssm",
-        skip_connections: bool = False,
-        *,
-        key: PRNGKeyArray,
-    ) -> None:
-        vocab_key, s4_key = jax.random.split(key, 2)
-        self.vocab_embedding = eqx.nn.Embedding(input_size, hidden_size, key=vocab_key)
-        self.proj_in = eqx.nn.Linear(hidden_size, hidden_size, key=key)
-        self.proj_out = eqx.nn.Linear(hidden_size, output_size, key=key)
-
-        block_keys = jax.random.split(s4_key, num_layers)
-
-        def get_block(key: PRNGKeyArray) -> BaseSSMBlock:
-            match block_type:
-                case "ssm":
-                    return SSMBlock(hidden_size, key=key)
-                case "diag":
-                    return DiagSSMBlock(hidden_size, key=key)
-                case _:
-                    raise ValueError(f"Unknown block type: {block_type}")
-
-        self.blocks = [get_block(block_keys[i]) for i in range(num_layers)]
-        self.skip_connections = skip_connections
-        self.num_layers = num_layers
-        self.hidden_size = hidden_size
-
-    def __call__(self, hs: list[Array], x: Array) -> tuple[list[Array], Array]:
-        new_hs = []
-        for i, block in enumerate(self.blocks):
-            h = block.forward(hs[i], x)
-            new_hs.append(h)
-            xh = jax.nn.gelu(h)
-            x = xh + x if self.skip_connections else xh
-        y = self.proj_out(x)
-        return new_hs, y
-
-    def _embed_input(self, x: Array) -> Array:
-        """U is the input to the S4 cell."""
-        embedded = self.vocab_embedding(x)
-        return jax.nn.gelu(self.proj_in(embedded))
-
-    def predict_sequence(self, x_seq: Array) -> Array:
-        x_emb = jax.vmap(self._embed_input)(x_seq)
-        hs = [jnp.zeros(self.hidden_size) for _ in range(self.num_layers)]
-
-        def step(hs: list[Array], x: Array) -> tuple[list[Array], Array]:
-            hs, y = self(hs, x)
-            return hs, y
-
-        _, y_seq = jax.lax.scan(step, hs, x_emb)
-        return y_seq