xax 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.9"
+__version__ = "0.1.11"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -43,9 +43,16 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
+    "cross_entropy",
     "cast_norm_type",
     "get_norm",
     "is_master",
+    "DiagSSMBlock",
+    "DiscreteTimeS4",
+    "S4",
+    "S4Layer",
+    "S6Layer",
+    "SSMBlock",
     "BaseLauncher",
     "CliLauncher",
     "SingleProcessLauncher",
@@ -196,9 +203,16 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
+    "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.norm",
     "get_norm": "nn.norm",
     "is_master": "nn.parallel",
+    "DiagSSMBlock": "nn.ssm",
+    "DiscreteTimeS4": "nn.ssm",
+    "S4": "nn.ssm",
+    "S4Layer": "nn.ssm",
+    "S6Layer": "nn.ssm",
+    "SSMBlock": "nn.ssm",
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
     "SingleProcessLauncher": "task.launchers.single_process",
@@ -351,8 +365,10 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_projected_gravity_vector_from_quat,
         quat_to_euler,
     )
+    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.task.base import RawConfigType
     from xax.task.launchers.base import BaseLauncher
     from xax.task.launchers.cli import CliLauncher

xax/nn/losses.py

@@ -0,0 +1,9 @@
+"""Defines some common loss functions."""
+
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def cross_entropy(y: Array, pred_y: Array, axis: int = 1) -> Array:
+    pred_y = jnp.take_along_axis(pred_y, jnp.expand_dims(y, axis), axis=axis)
+    return -jnp.mean(pred_y)

xax/nn/norm.py

@@ -3,6 +3,7 @@
 from typing import Literal, cast, get_args
 
 import jax.numpy as jnp
+from jaxtyping import Array
 
 NormType = Literal["l1", "l2"]
 
@@ -13,7 +14,7 @@ def cast_norm_type(norm: str) -> NormType:
     return cast(NormType, norm)
 
 
-def get_norm(x: jnp.ndarray, norm: NormType) -> jnp.ndarray:
+def get_norm(x: Array, norm: NormType) -> Array:
     match norm:
         case "l1":
             return jnp.abs(x)

xax/nn/ssm.py

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

xax/task/mixins/train.py

@@ -29,6 +29,7 @@ from typing import (
 
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 import numpy as np
 import optax
 from jaxtyping import Array, PRNGKeyArray, PyTree
@@ -56,6 +57,7 @@ from xax.utils.experiments import (
 from xax.utils.jax import jit as xax_jit
 from xax.utils.logging import LOG_STATUS
 from xax.utils.text import highlight_exception_message, show_info
+from xax.utils.types.frozen_dict import FrozenDict
 
 logger = logging.getLogger(__name__)
 
@@ -161,6 +163,7 @@ class TrainConfig(
     max_steps: int | None = field(None, help="Maximum number of steps to run")
     step_kind: str = field("step", help=f"How to measure a step; one of [{', '.join(get_args(StepKind))}]")
     random_seed: int = field(1337, help="Random seed for the task")
+    global_grad_clip: float = field(value=10.0, help="The maximum gradient norm to clip to.")
 
 
 Config = TypeVar("Config", bound=TrainConfig)
@@ -215,7 +218,7 @@ 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, state: State) -> None:
+    def log_train_step(self, 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
@@ -224,10 +227,11 @@ class TrainMixin(
         Args:
             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, state: State) -> None:
+    def log_valid_step(self, 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
@@ -236,6 +240,7 @@ class TrainMixin(
         Args:
             batch: The batch from the dataloader.
             output: The model output.
+            metrics: The metrics for the current batch.
             state: The current training state.
         """
 
@@ -246,18 +251,23 @@ class TrainMixin(
             for k, v in d.items():
                 self.logger.log_scalar(k, v, namespace=ns)
 
-    def log_step(self, batch: Batch, output: Output, loss: Array, state: State) -> None:
+    def log_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
         phase = state.phase
 
-        self.logger.log_scalar("loss", loss, namespace="loss")
+        for k, v in metrics.items():
+            if v.size == 1:
+                self.logger.log_scalar(k, v.item())
+            else:
+                self.logger.log_histogram(k, v)
+
         self.log_state_timers(state)
 
         # Delegate to the appropriate logging function based on the phase.
         match phase:
             case "train":
-                self.log_train_step(batch, output, state)
+                self.log_train_step(batch, output, metrics, state)
             case "valid":
-                self.log_valid_step(batch, output, state)
+                self.log_valid_step(batch, output, metrics, state)
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
@@ -364,32 +374,90 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
+    def compute_metrics(
+        self,
+        model: PyTree,
+        batch: Batch,
+        output: Output,
+        loss: Array,
+        state: State,
+    ) -> dict[str, Array]:
+        """Computes the metrics for the current batch.
+
+        Args:
+            model: The current model.
+            batch: The current minibatch of samples.
+            output: The output from the model.
+            loss: The loss for the current batch.
+            state: The current training state.
+
+        Returns:
+            A dictionary of metrics.
+        """
+        return {
+            "loss": loss,
+        }
+
+    @xax_jit(static_argnames=["self", "model_static"])
     def get_output_and_loss(
         self,
-        model_static: PyTree,
         model_arr: PyTree,
+        model_static: PyTree,
         batch: Batch,
         state: State,
-    ) -> tuple[Array, Output]:
+    ) -> tuple[Array, tuple[Output, dict[str, Array]]]:
         model = eqx.combine(model_arr, model_static)
         output = self.get_output(model, batch, state)
         loss = self.compute_loss(model, batch, output, state)
-        return loss, output
+        metrics = self.compute_metrics(model, batch, output, loss, state)
+        return loss, (output, metrics)
 
     def update(
         self,
-        model_static: PyTree,
         model_arr: PyTree,
+        model_static: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
         state: State,
-    ) -> tuple[Array, PyTree, optax.OptState, Output]:
-        grad_fn = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)
-        (loss, output), grads = grad_fn(model_static, model_arr, batch, state)
-        updates, opt_state = optimizer.update(grads, opt_state, model_arr)
-        model_arr = eqx.apply_updates(model_arr, updates)
-        return loss, model_arr, opt_state, output
+    ) -> tuple[PyTree, optax.OptState, Output, dict[str, Array]]:
+        grad_fn = jax.grad(self.get_output_and_loss, argnums=0, has_aux=True)
+        grad_fn = xax_jit(static_argnums=[1])(grad_fn)
+        grads, (output, metrics) = grad_fn(model_arr, model_static, batch, state)
+        model_arr, opt_state, grad_metrics = self.apply_gradients_with_clipping(model_arr, grads, optimizer, opt_state)
+        return model_arr, opt_state, output, metrics | grad_metrics
+
+    @xax_jit(static_argnames=["self", "optimizer"])
+    def apply_gradients_with_clipping(
+        self,
+        model_arr: PyTree,
+        grads: PyTree,
+        optimizer: optax.GradientTransformation,
+        opt_state: optax.OptState,
+    ) -> tuple[PyTree, optax.OptState, dict[str, Array]]:
+        grad_norm = optax.global_norm(grads)
+        grad_metrics = {"grad_norm": grad_norm}
+
+        def apply(grads: PyTree, grad_norm: Array) -> tuple[PyTree, optax.OptState]:
+            # Clip the global gradient norm to some desired range.
+            grad_factor = self.config.global_grad_clip / jnp.maximum(grad_norm, 1e-6)
+            grads = jax.tree.map(lambda x: x * grad_factor, grads)
+
+            # Apply the gradient updates.
+            updates, new_opt_state = optimizer.update(grads, opt_state, model_arr)
+            new_model_arr = eqx.apply_updates(model_arr, updates)
+            return new_model_arr, new_opt_state
+
+        # Don't apply updates if the gradient is NaN or Inf.
+        new_model_arr, new_opt_state = jax.lax.cond(
+            jnp.isnan(grad_norm) | jnp.isinf(grad_norm),
+            lambda *_: (model_arr, opt_state),
+            apply,
+            grads,
+            grad_norm,
+        )
+
+        return new_model_arr, new_opt_state, grad_metrics
 
     def get_size_of_batch(self, batch: Batch) -> int | None:
         """Gets the batch size for the current batch.
@@ -469,25 +537,26 @@ class TrainMixin(
     @xax_jit(static_argnames=["self", "model_static", "optimizer"])
     def train_step(
         self,
-        model_static: PyTree,
         model_arr: PyTree,
+        model_static: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
         state: State,
-    ) -> tuple[PyTree, optax.OptState, Array, Output]:
-        loss, model_arr, opt_state, output = self.update(model_static, model_arr, optimizer, opt_state, batch, state)
-        return model_arr, opt_state, loss, output
+    ) -> tuple[PyTree, optax.OptState, Output, FrozenDict[str, Array]]:
+        model_arr, opt_state, output, metrics = self.update(model_arr, model_static, optimizer, opt_state, batch, state)
+        return model_arr, opt_state, output, FrozenDict(metrics)
 
     @xax_jit(static_argnames=["self", "model_static"])
     def val_step(
         self,
-        model_static: PyTree,
         model_arr: PyTree,
+        model_static: PyTree,
         batch: Batch,
         state: State,
-    ) -> tuple[Array, Output]:
-        return self.get_output_and_loss(model_static, model_arr, batch, state)
+    ) -> tuple[Output, FrozenDict[str, Array]]:
+        _, (output, metrics) = self.get_output_and_loss(model_arr, model_static, batch, state)
+        return output, FrozenDict(metrics)
 
     def train_loop(
         self,
@@ -509,8 +578,8 @@ class TrainMixin(
                     num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
                 )
 
-                loss, output = self.val_step(model_static, model_arr, valid_batch, state)
-                self.log_step(valid_batch, output, loss, state)
+                output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
+                self.log_step(valid_batch, output, metrics, state)
 
             state = self.on_step_start(state)
             train_batch = next(train_pf)
@@ -520,15 +589,15 @@ class TrainMixin(
                 num_samples=state.num_samples + (self.get_size_of_batch(train_batch) or 0),
             )
 
-            model_arr, opt_state, loss, output = self.train_step(
-                model_static=model_static,
+            model_arr, opt_state, output, metrics = self.train_step(
                 model_arr=model_arr,
+                model_static=model_static,
                 optimizer=optimizer,
                 opt_state=opt_state,
                 batch=train_batch,
                 state=state,
             )
-            self.log_step(train_batch, output, loss, state)
+            self.log_step(train_batch, output, metrics, state)
 
             state = self.on_step_end(state)
 

xax/utils/pytree.py

@@ -31,7 +31,7 @@ def slice_array(x: Array, start: Array, slice_length: int) -> Array:
 
 def slice_pytree(pytree: PyTree, start: Array, slice_length: int) -> PyTree:
     """Get a slice of a pytree."""
-    return jax.tree_util.tree_map(lambda x: slice_array(x, start, slice_length), pytree)
+    return jax.tree.map(lambda x: slice_array(x, start, slice_length), pytree)
 
 
 def flatten_array(x: Array, flatten_size: int) -> Array:
@@ -43,14 +43,14 @@ def flatten_array(x: Array, flatten_size: int) -> Array:
 
 def flatten_pytree(pytree: PyTree, flatten_size: int) -> PyTree:
     """Flatten a pytree into a (flatten_size, ...) pytree."""
-    return jax.tree_util.tree_map(lambda x: flatten_array(x, flatten_size), pytree)
+    return jax.tree.map(lambda x: flatten_array(x, flatten_size), pytree)
 
 
 def pytree_has_nans(pytree: PyTree) -> Array:
     """Check if a pytree has any NaNs."""
     has_nans = jax.tree_util.tree_reduce(
         lambda a, b: jnp.logical_or(a, b),
-        jax.tree_util.tree_map(lambda x: jnp.any(jnp.isnan(x)), pytree),
+        jax.tree.map(lambda x: jnp.any(jnp.isnan(x)), pytree),
     )
     return has_nans
 
@@ -58,13 +58,13 @@ def pytree_has_nans(pytree: PyTree) -> Array:
 def update_pytree(cond: Array, new: PyTree, original: PyTree) -> PyTree:
     """Update a pytree based on a condition."""
     # Tricky, need use tree_map because where expects array leafs.
-    return jax.tree_util.tree_map(lambda x, y: jnp.where(cond, x, y), new, original)
+    return jax.tree.map(lambda x, y: jnp.where(cond, x, y), new, original)
 
 
 def compute_nan_ratio(pytree: PyTree) -> Array:
     """Computes the ratio of NaNs vs non-NaNs in a given PyTree."""
-    nan_counts = jax.tree_util.tree_map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
-    total_counts = jax.tree_util.tree_map(lambda x: x.size, pytree)
+    nan_counts = jax.tree.map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
+    total_counts = jax.tree.map(lambda x: x.size, pytree)
 
     total_nans = jax.tree_util.tree_reduce(lambda a, b: a + b, nan_counts, 0)
     total_elements = jax.tree_util.tree_reduce(lambda a, b: a + b, total_counts, 0)
@@ -118,7 +118,7 @@ def reshuffle_pytree(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArr
         # Reshape back to the original shape
         return permuted.reshape(orig_shape)
 
-    return jax.tree_util.tree_map(permute_array, data)
+    return jax.tree.map(permute_array, data)
 
 
 def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArray) -> PyTree:
@@ -133,7 +133,7 @@ def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], r
             return x[tuple(idx_grids)]
         return x
 
-    return jax.tree_util.tree_map(permute_array, data)
+    return jax.tree.map(permute_array, data)
 
 
 TransposeResult = tuple[PyTree, tuple[int, ...], tuple[int, ...]]
@@ -215,7 +215,7 @@ def reshuffle_pytree_along_dims(
                 transpose_info[path] = (transpose_order, original_shape)
         return x
 
-    jax.tree_util.tree_map_with_path(prepare_for_shuffle, data)
+    jax.tree.map_with_path(prepare_for_shuffle, data)
 
     # Create a transposed pytree
     def get_transposed(path: PathType, x: PyTree) -> PyTree:
@@ -223,7 +223,7 @@ def reshuffle_pytree_along_dims(
             return transposed_data[path]
         return x
 
-    transposed_pytree = jax.tree_util.tree_map_with_path(get_transposed, data)
+    transposed_pytree = jax.tree.map_with_path(get_transposed, data)
 
     # Reshuffle the transposed pytree along the leading dimensions
     reshuffled_transposed = reshuffle_pytree(transposed_pytree, shape_dims, rng)
@@ -235,4 +235,4 @@ def reshuffle_pytree_along_dims(
             return transpose_back(x, transpose_order, original_shape)
         return x
 
-    return jax.tree_util.tree_map_with_path(restore_transpose, reshuffled_transposed)
+    return jax.tree.map_with_path(restore_transpose, reshuffled_transposed)

{xax-0.1.9.dist-info → xax-0.1.11.dist-info}/METADATA

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

{xax-0.1.9.dist-info → xax-0.1.11.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=_xb60-jl7arZEleSwUw4ElPaq4MzD24_ZYQrnWO5_cs,13391
+xax/__init__.py,sha256=2JdSxsZphJJFVMGBVXNc0hP2p0FVOu5y7xSgPRNeyNY,13835
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -11,8 +11,10 @@ xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
 xax/nn/export.py,sha256=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
 xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
-xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
+xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
+xax/nn/norm.py,sha256=WgZ3QCrUnf-YecwhEtVPcr99fKK3ECl_UeiAs2uv7oo,564
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
+xax/nn/ssm.py,sha256=eFeGkV1pkVGc0vNrQbykCbFnlPXQqsqVA_JVzLBHD28,9865
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/base.py,sha256=E4l1yCrAkM2TVTbVYrmk6BoVHMkbD4IYsTT921XOyi0,7760
 xax/task/logger.py,sha256=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
@@ -39,7 +41,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=JbrSiBqpgOrdDanNYuAzzh2radPrXOVrHYA6VcxjIzY,23248
+xax/task/mixins/train.py,sha256=lgLHiHQtnDK0XS3SwHTYZtDv5CTbPRN1-p_K9KiIpHQ,26000
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
@@ -48,7 +50,7 @@ xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=7GjQoPc_ZSZt3QS_9qXoBWl1jfMp1qZa7aViQoWJ0OQ,8864
+xax/utils/pytree.py,sha256=VFWhT0MQ99KjQyEYM6NFbqYq4_hOZwB23uhowMB4U34,8754
 xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,17058
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -56,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.9.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.9.dist-info/METADATA,sha256=Ou8KmYWWNxgo_9ZAU2KLaeGeXAxd6b9qJ95ky4HRm-o,1877
-xax-0.1.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.9.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.9.dist-info/RECORD,,
+xax-0.1.11.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.11.dist-info/METADATA,sha256=qDhn5EGxdiuEe5gQUZiBC430sXhJOPRWboTvsh2onxs,1878
+xax-0.1.11.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.11.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.11.dist-info/RECORD,,