xax 0.2.12__tar.gz → 0.2.13__tar.gz

{xax-0.2.12/xax.egg-info → xax-0.2.13}/PKG-INFO

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

{xax-0.2.12 → xax-0.2.13}/xax/__init__.py

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

{xax-0.2.12 → xax-0.2.13}/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-0.2.12 → xax-0.2.13}/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-0.2.12 → xax-0.2.13}/xax/task/mixins/train.py

@@ -728,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(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,
@@ -754,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.12 → xax-0.2.13/xax.egg-info}/PKG-INFO

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