xax 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.11"
+__version__ = "0.2.13"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/core/state.py

@@ -1,5 +1,6 @@
 """Defines a dataclass for keeping track of the current training state."""
 
+import time
 from dataclasses import dataclass
 from typing import Literal, NotRequired, TypedDict, Unpack, cast
 
@@ -28,6 +29,7 @@ class StateDict(TypedDict, total=False):
     num_samples: NotRequired[int | Array]
     num_valid_steps: NotRequired[int | Array]
     num_valid_samples: NotRequired[int | Array]
+    start_time_s: NotRequired[float | Array]
     elapsed_time_s: NotRequired[float | Array]
     valid_elapsed_time_s: NotRequired[float | Array]
     phase: NotRequired[Phase]
@@ -57,13 +59,17 @@ class State:
         return self._float32_arr[1]
 
     @property
-    def elapsed_time_s(self) -> Array:
+    def start_time_s(self) -> Array:
         return self._float32_arr[2]
 
     @property
-    def valid_elapsed_time_s(self) -> Array:
+    def elapsed_time_s(self) -> Array:
         return self._float32_arr[3]
 
+    @property
+    def valid_elapsed_time_s(self) -> Array:
+        return self._float32_arr[4]
+
     @property
     def phase(self) -> Phase:
         return _int_to_phase(self._int32_arr[2].item())
@@ -72,7 +78,7 @@ class State:
     def init_state(cls) -> "State":
         return cls(
             _int32_arr=jnp.array([0, 0, 0], dtype=jnp.int32),
-            _float32_arr=jnp.array([0.0, 0.0, 0.0, 0.0], dtype=jnp.float32),
+            _float32_arr=jnp.array([0.0, 0.0, time.time(), 0.0, 0.0], dtype=jnp.float32),
         )
 
     @property
@@ -98,10 +104,12 @@ class State:
         if "num_valid_samples" in kwargs:
             float32_arr = float32_arr.at[1].set(kwargs["num_valid_samples"])
 
+        if "start_time_s" in kwargs:
+            float32_arr = float32_arr.at[2].set(kwargs["start_time_s"])
         if "elapsed_time_s" in kwargs:
-            float32_arr = float32_arr.at[2].set(kwargs["elapsed_time_s"])
+            float32_arr = float32_arr.at[3].set(kwargs["elapsed_time_s"])
         if "valid_elapsed_time_s" in kwargs:
-            float32_arr = float32_arr.at[3].set(kwargs["valid_elapsed_time_s"])
+            float32_arr = float32_arr.at[4].set(kwargs["valid_elapsed_time_s"])
 
         return State(
             _int32_arr=int32_arr,
@@ -110,12 +118,13 @@ class State:
 
     def to_dict(self) -> dict[str, int | float | str]:
         return {
-            "num_steps": int(self.num_steps),
-            "num_valid_steps": int(self.num_valid_steps),
-            "num_samples": int(self.num_samples),
-            "num_valid_samples": int(self.num_valid_samples),
-            "elapsed_time_s": float(self.elapsed_time_s),
-            "valid_elapsed_time_s": float(self.valid_elapsed_time_s),
+            "num_steps": int(self.num_steps.item()),
+            "num_valid_steps": int(self.num_valid_steps.item()),
+            "num_samples": int(self.num_samples.item()),
+            "num_valid_samples": int(self.num_valid_samples.item()),
+            "start_time_s": float(self.start_time_s.item()),
+            "elapsed_time_s": float(self.elapsed_time_s.item()),
+            "valid_elapsed_time_s": float(self.valid_elapsed_time_s.item()),
             "phase": str(self.phase),
         }
 
@@ -137,6 +146,7 @@ class State:
             [
                 d.get("num_samples", 0),
                 d.get("num_valid_samples", 0),
+                d.get("start_time_s", time.time()),
                 d.get("elapsed_time_s", 0.0),
                 d.get("valid_elapsed_time_s", 0.0),
             ],

xax/task/loggers/tensorboard.py

@@ -160,7 +160,7 @@ class TensorboardLogger(LoggerImpl):
         writer = self.get_writer(line.state.phase)
 
         global_step = line.state.num_steps.item()
-        walltime = line.state.elapsed_time_s.item()
+        walltime = line.state.start_time_s.item() + line.state.elapsed_time_s.item()
 
         for namespace, scalars in line.scalars.items():
             for scalar_key, scalar_value in scalars.items():

xax/task/mixins/train.py

@@ -120,7 +120,7 @@ class ValidStepTimer:
         self.last_valid_time = state.elapsed_time_s.item()
         self.last_valid_step = state.num_steps.item()
 
-    def is_valid_step(self, state: State) -> bool:
+    def __call__(self, state: State) -> bool:
         if state.num_steps < self.valid_first_n_steps and state.num_valid_steps < self.valid_first_n_steps:
             return True
 
@@ -130,15 +130,18 @@ class ValidStepTimer:
 
         # Step-based validation.
         valid_every_n_steps = self.valid_every_n_steps
-        if valid_every_n_steps is not None and state.num_steps >= valid_every_n_steps + self.last_valid_step:
+        if valid_every_n_steps is not None and (
+            state.num_steps >= valid_every_n_steps + self.last_valid_step
+            or state.num_valid_steps >= valid_every_n_steps + self.last_valid_step
+        ):
             self._reset(state)
             return True
 
         # Time-based validation.
         valid_every_n_seconds = self.valid_every_n_seconds
-        if (
-            valid_every_n_seconds is not None
-            and state.elapsed_time_s.item() - self.last_valid_time >= valid_every_n_seconds
+        if valid_every_n_seconds is not None and (
+            state.elapsed_time_s.item() - self.last_valid_time >= valid_every_n_seconds
+            or state.valid_elapsed_time_s.item() - self.last_valid_time >= valid_every_n_seconds
         ):
             self._reset(state)
             return True
@@ -146,7 +149,10 @@ class ValidStepTimer:
         # Time-based validation for first validation step.
         if self.first_valid_step_flag:
             valid_first_n_seconds = self.valid_first_n_seconds
-            if valid_first_n_seconds is not None and state.elapsed_time_s.item() >= valid_first_n_seconds:
+            if valid_first_n_seconds is not None and (
+                state.elapsed_time_s.item() >= valid_first_n_seconds
+                or state.valid_elapsed_time_s.item() >= valid_first_n_seconds
+            ):
                 self._reset(state)
                 self.first_valid_step_flag = False
                 return True
@@ -722,21 +728,27 @@ class TrainMixin(
         model_arr, model_static = eqx.partition(model, self.model_partition_fn)
 
         while not self.is_training_over(state):
-            if self.valid_step_timer.is_valid_step(state):
+            valid_step = self.valid_step_timer(state)
+
+            if valid_step:
                 with ContextTimer() as timer:
+                    state = state.replace(phase="valid")
                     valid_batch = next(valid_pf)
                     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_valid_steps=state.num_valid_steps + 1,
+                        num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
+                    )
+
                 state = state.replace(
-                    phase="valid",
-                    num_valid_steps=state.num_valid_steps + 1,
-                    num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
                     valid_elapsed_time_s=state.valid_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)
                 model_arr, opt_state, output, metrics = self.train_step(
                     model_arr=model_arr,
@@ -748,15 +760,17 @@ class TrainMixin(
                 )
                 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),
+                )
+
+                state = self.on_step_end(state)
+
             state = state.replace(
-                phase="train",
-                num_steps=state.num_steps + 1,
-                num_samples=state.num_samples + (self.get_size_of_batch(train_batch) or 0),
                 elapsed_time_s=state.elapsed_time_s + timer.elapsed_time,
             )
 
-            state = self.on_step_end(state)
-
             if self.should_checkpoint(state):
                 model = eqx.combine(model_arr, model_static)
                 self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)

{xax-0.2.11.dist-info → xax-0.2.13.dist-info}/METADATA

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

{xax-0.2.11.dist-info → xax-0.2.13.dist-info}/RECORD

@@ -1,10 +1,10 @@
-xax/__init__.py,sha256=S4p0bL4JmuLyhFkGHlXlugXk-ckbnWtSw1_6r9E0qrI,15510
+xax/__init__.py,sha256=33wIwGeXDFReg2ZnFqUHfSybj5cKyMqnI8ncj8-9yVg,15510
 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
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
-xax/core/state.py,sha256=yO25lMoLCUTJlHyLzQxlDbsHC_GZ3HkrKAq5huA7AkU,4552
+xax/core/state.py,sha256=KsNMnM_RgsZ2Ntc2pp4Fi6zG4rZb_89-kqmyGxDvyRg,4974
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=JZuSApD4bL0UK5W1nrQtucWYvNWUha07J6LTLk_RX-Y,4910
@@ -29,7 +29,7 @@ xax/task/loggers/callback.py,sha256=zQuV1xCvz47Q3UQqP1D5mBhbVzptvmPR_7hX25vqSk0,
 xax/task/loggers/json.py,sha256=6A5wL7kspsXnpPhI_vu0scgd2Z2-WLhw4gbBFm7eZMM,4377
 xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,1407
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
-xax/task/loggers/tensorboard.py,sha256=sdsA8GjZG5JQpoAxNDRr_bGvqN8Olgj_almZBb2K5F8,8850
+xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
 xax/task/mixins/checkpointing.py,sha256=zqospBFnTbGt_iriiduVfXazINPbzWpwmIs91KAniMY,10147
@@ -41,7 +41,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 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=PUVN2OsJpQppIzb4yaULT-C-0ocr1aGbPY-LrNJ2AVY,31322
+xax/task/mixins/train.py,sha256=_QoxSDMW6nmpH82Un2LDsVIBg9KIx8npRwSjY4TEGYA,31830
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=bj8BftSHT3fFzfiJ0Co0WvqWo0rUS8kQnQYpVvH8FTM,29942
@@ -58,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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.11.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.11.dist-info/METADATA,sha256=MLwHl-vblIYvbpUZ5ylMDjwejKLNOnJK_55JwVNPVH8,1880
-xax-0.2.11.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-xax-0.2.11.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.11.dist-info/RECORD,,
+xax-0.2.13.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.13.dist-info/METADATA,sha256=-foHRw3ph7yxBmMmjO_oqZqwvdEROYTH4Drc9P58ujI,1880
+xax-0.2.13.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+xax-0.2.13.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.13.dist-info/RECORD,,