xax 0.1.9__py3-none-any.whl → 0.1.10__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.10"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/task/mixins/train.py

@@ -56,6 +56,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__)
 
@@ -215,7 +216,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 +225,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 +238,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 +249,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 +372,59 @@ 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, FrozenDict[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, FrozenDict(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)
+    ) -> tuple[PyTree, optax.OptState, Output, FrozenDict[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)
         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
+        return model_arr, opt_state, output, metrics
 
     def get_size_of_batch(self, batch: Batch) -> int | None:
         """Gets the batch size for the current batch.
@@ -469,25 +504,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, 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, metrics
 
     def train_loop(
         self,
@@ -509,8 +545,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 +556,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.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.9
+Version: 0.1.10
 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.10.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=_xb60-jl7arZEleSwUw4ElPaq4MzD24_ZYQrnWO5_cs,13391
+xax/__init__.py,sha256=bvOBMlEVA46I7ILGfk5AbpwpcdTAjw-4vWI7ci7L7-g,13392
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -39,7 +39,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=jAzc9RD25DbhekvItzsRQQrK9aEwtA_sXy0m2Hfkuxo,24594
 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 +48,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 +56,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.10.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.10.dist-info/METADATA,sha256=kJ1lxZ6cWrtJ5R-adTorzEE_1l0VRJ67xfuBjYXG9Vo,1878
+xax-0.1.10.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.10.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.10.dist-info/RECORD,,