xax 0.3.8__py3-none-any.whl → 0.3.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.3.8"
+__version__ = "0.3.10"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/task/mixins/data_loader.py

@@ -110,7 +110,12 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             "or `get_data_iterator` to return an iterator for the given dataset."
         )
 
-    def get_dataloader(self, dataset: Dataset[T, Tc_co], phase: Phase) -> Dataloader[T, Tc_co]:
+    def get_dataloader(
+        self,
+        dataset: Dataset[T, Tc_co],
+        phase: Phase,
+        prefetch_factor: int | None = None,
+    ) -> Dataloader[T, Tc_co]:
         debugging = self.config.debug_dataloader
         if debugging:
             logger.warning("Parallel dataloaders disabled in debugging mode")
@@ -135,7 +140,7 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             dataset=dataset,
             batch_size=self.config.batch_size,
             num_workers=0 if debugging else cfg.num_workers,
-            prefetch_factor=cfg.prefetch_factor,
+            prefetch_factor=cfg.prefetch_factor if prefetch_factor is None else prefetch_factor,
             mp_manager=self.multiprocessing_manager,
             dataloader_worker_init_fn=self.dataloader_worker_init_fn,
             collate_worker_init_fn=self.collate_worker_init_fn,

xax/task/mixins/train.py

@@ -60,7 +60,7 @@ from xax.utils.experiments import (
     get_state_file_string,
     get_training_code,
 )
-from xax.utils.jax import jit as xax_jit
+from xax.utils.jax import jit as xax_jit, scan as xax_scan
 from xax.utils.logging import LOG_PING, LOG_STATUS
 from xax.utils.pytree import get_pytree_param_count
 from xax.utils.text import highlight_exception_message, show_info
@@ -175,6 +175,7 @@ class TrainConfig(
     valid_first_n_seconds: float | None = field(60.0, help="Run first validation after N seconds")
     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))}]")
+    updates_per_step: int = field(1, help="Number of updates to perform per step")
     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.")
 
@@ -597,6 +598,7 @@ class TrainMixin(
         metrics = self.compute_metrics(model, batch, output, loss, state)
         return loss, (output, metrics)
 
+    @xax_jit(static_argnames=["self", "model_static", "optimizer"], jit_level=3)
     def update(
         self,
         model_arr: PyTree,
@@ -609,44 +611,9 @@ class TrainMixin(
         grad_fn = jax.grad(self.get_output_and_loss, argnums=0, has_aux=True)
         grad_fn = xax_jit(static_argnums=[1], jit_level=3)(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"], jit_level=3)
-    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 gradients based on global norm, similar to optax.clip_by_global_norm
-            trigger = jnp.squeeze(grad_norm < self.config.global_grad_clip)
-
-            def clip_fn(t: Array) -> Array:
-                return jax.lax.select(trigger, t, (t / grad_norm.astype(t.dtype)) * self.config.global_grad_clip)
-
-            grads = jax.tree.map(clip_fn, 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
+        updates, opt_state = optimizer.update(grads, opt_state, model_arr)
+        model_arr = eqx.apply_updates(model_arr, updates)
+        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.
@@ -729,11 +696,36 @@ class TrainMixin(
         model_static: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
-        batch: Batch,
+        batches: Batch,
         state: State,
     ) -> 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)
+        def update_fn(
+            carry: tuple[PyTree, optax.OptState],
+            batch: Batch,
+        ) -> tuple[tuple[PyTree, optax.OptState], tuple[Output, FrozenDict[str, Array]]]:
+            model_arr, opt_state = carry
+            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))
+
+        (model_arr, opt_state), (output, metrics) = xax_scan(
+            update_fn,
+            (model_arr, opt_state),
+            batches,
+            jit_level=3,
+        )
+
+        # Only get the final output and metrics.
+        output = jax.tree.map(lambda x: x[-1], output)
+        metrics = jax.tree.map(lambda x: x[-1], metrics)
+
+        return model_arr, opt_state, output, metrics
 
     @xax_jit(static_argnames=["self", "model_static"], jit_level=3)
     def val_step(
@@ -775,40 +767,36 @@ class TrainMixin(
                     output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
                     self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
 
-                    state = state.replace(
-                        num_steps=state.num_steps + 1,
-                        num_samples=state.num_samples + (self.get_size_of_batch(valid_batch) or 0),
-                    )
-
                 state = state.replace(
+                    num_steps=state.num_steps + 1,
+                    num_samples=state.num_samples + (self.get_size_of_batch(valid_batch) or 0),
                     elapsed_time_s=state.elapsed_time_s + timer.elapsed_time,
                 )
 
             with ContextTimer() as timer:
                 state = self.on_step_start(state)
                 state = state.replace(phase="train")
-                train_batch = next(train_pf)
+                train_batches = list(itertools.islice(train_pf, self.config.updates_per_step))
                 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,
+                    batches=jax.tree.map(lambda *xs: jnp.stack(xs, axis=0), *train_batches),
                     state=state,
                 )
-                self.log_step(eqx.combine(model_arr, model_static), train_batch, output, metrics, state)
-
-                state = state.replace(
-                    num_steps=state.num_steps + 1,
-                    num_samples=state.num_samples + (self.get_size_of_batch(train_batch) or 0),
-                )
-
+                self.log_step(eqx.combine(model_arr, model_static), train_batches[-1], output, metrics, state)
                 state = self.on_step_end(state)
 
             state = state.replace(
+                num_steps=state.num_steps + 1,
+                num_samples=state.num_samples + (self.get_size_of_batch(train_batches[-1]) or 0),
                 elapsed_time_s=state.elapsed_time_s + timer.elapsed_time,
             )
 
+            if state.num_steps <= 3:
+                logger.log(LOG_PING, "Step %d took %.2f second", state.num_steps, timer.elapsed_time)
+
             if self.should_checkpoint(state):
                 model = eqx.combine(model_arr, model_static)
                 self.save_checkpoint(models=[model], optimizers=[optimizer], opt_states=[opt_state], state=state)
@@ -827,7 +815,7 @@ class TrainMixin(
             pass
 
         train_ds = self.get_dataset("train")
-        train_dl = self.get_dataloader(train_ds, "train")
+        train_dl = self.get_dataloader(train_ds, "train", prefetch_factor=self.config.updates_per_step + 1)
         train_pf = self.get_prefetcher(train_dl)
 
         try:

xax/utils/pytree.py

@@ -1,6 +1,6 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
-from typing import TypeVar
+from typing import Mapping, Sequence, TypeVar
 
 import chex
 import equinox as eqx
@@ -258,11 +258,18 @@ def tuple_insert(t: tuple[T, ...], index: int, value: T) -> tuple[T, ...]:
 def get_pytree_mapping(pytree: PyTree) -> dict[str, Array]:
     leaves: dict[str, Array] = {}
 
+    def _get_str(thing: PyTree) -> str:
+        if isinstance(thing, str):
+            return thing
+        if isinstance(thing, Sequence):
+            return "/".join(_get_str(x) for x in thing)
+        if isinstance(thing, Mapping):
+            return "/".join(f"{_get_str(k)}:{_get_str(v)}" for k, v in thing.items())
+        return str(thing)
+
     def _get_leaf(path: tuple, x: PyTree) -> None:
         if isinstance(x, jnp.ndarray):
-            # Convert path tuple to string, e.g. (1, 'a', 2) -> '1/a/2'
-            path_str = "/".join(str(p) for p in path)
-            leaves[path_str] = x
+            leaves[_get_str(path)] = x
 
     jax.tree.map_with_path(_get_leaf, pytree)
     return leaves

{xax-0.3.8.dist-info → xax-0.3.10.dist-info}/METADATA

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

{xax-0.3.8.dist-info → xax-0.3.10.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=5NGaVm9X36LhG-Tl1hc7Lk1SmnTZvyu8G1iFDixpqLc,16665
+xax/__init__.py,sha256=lSwyrPTof_BZ-pyPNhNICJnCZMN9i2sJ-Ii3S_vY_28,16666
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
@@ -37,13 +37,13 @@ xax/task/mixins/artifacts.py,sha256=R-y3p7__zJHlHDqwDVAZysg2ZmebCJbqAx_xGT2Xpd0,
 xax/task/mixins/checkpointing.py,sha256=v50IZ7j58DWmEu-_6Zh_02R5KUVGhrMkg5n-MYM_J4c,11484
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
-xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-goU,6599
+xax/task/mixins/data_loader.py,sha256=BKfOVWXR70vbyHMFlnlUiQQHXHH5zTj5WtmsymNCFB4,6722
 xax/task/mixins/gpu_stats.py,sha256=USOyhXldxbsrl6eCtoFKTWUm_lfeG0cUCkQNUpXRdtA,8880
 xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,2808
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 xax/task/mixins/runnable.py,sha256=pcLrYc_TycZUY9zZim05Skc2FWk3IZKFnu6p3UDMonM,1966
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=bjBoigTCjbq9H4hcqIO32irHBc9rC2zkgXrnGNI2RtI,33266
+xax/task/mixins/train.py,sha256=_kDpifLi1arSuT0ssFhBV0axpvLlQG3a97pohya0Eqc,32908
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=5k5hPYSaVjzoR_nm2Q3DAHMMYi3Bcp3N3PAQbwZq7Gg,29830
@@ -52,7 +52,7 @@ xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 xax/utils/logging.py,sha256=Kkyma_LJXqrN2HTQ214gRP_9ih3_bKk115MWC60lQWM,6656
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=cLZRSd5xc-DqcbRfWnBy87pAiUU5fT8U4CHoLi_i_v4,9642
+xax/utils/pytree.py,sha256=w8Ab2LmJdQ8e1FxKF0xWWaOak09Mhu44ZcOeUR6uGFA,9889
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
 xax/utils/text.py,sha256=xS02aSzdywl3KIaNSpKWcxdd37oYlUJtu9wIjkc1wVc,10654
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -60,9 +60,9 @@ 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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.3.8.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.8.dist-info/METADATA,sha256=d4UVJYHBKGAJTdC8G4IHt9kI44lbexOWIiZnkICd0pM,1246
-xax-0.3.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.8.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.8.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.8.dist-info/RECORD,,
+xax-0.3.10.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.10.dist-info/METADATA,sha256=oQMGYjsfYxMmw0A60qE15yda_G-0YG5RNl17tboR1f0,1247
+xax-0.3.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.10.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.10.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.10.dist-info/RECORD,,