xax 0.1.1__tar.gz → 0.1.2__tar.gz

{xax-0.1.1/xax.egg-info → xax-0.1.2}/PKG-INFO

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

{xax-0.1.1 → xax-0.1.2}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -66,6 +66,7 @@ __all__ = [
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
+    "ValidStepTimer",
     "Script",
     "ScriptConfig",
     "Config",
@@ -219,6 +220,7 @@ NAME_MAP: dict[str, str] = {
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
+    "ValidStepTimer": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
@@ -372,7 +374,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
     from xax.task.mixins.step_wrapper import StepContext
-    from xax.task.mixins.train import Batch, Output
+    from xax.task.mixins.train import Batch, Output, ValidStepTimer
     from xax.task.script import Script, ScriptConfig
     from xax.task.task import Config, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null

{xax-0.1.1 → xax-0.1.2}/xax/utils/experiments.py

@@ -109,8 +109,8 @@ class StateTimer:
 
     def step(self, state: State) -> None:
         cur_time = time.time()
-        self.step_timer.step(state.num_steps, cur_time)
-        self.sample_timer.step(state.num_samples, cur_time)
+        self.step_timer.step(state.num_steps if state.phase == "train" else state.num_valid_steps, cur_time)
+        self.sample_timer.step(state.num_samples if state.phase == "train" else state.num_valid_samples, cur_time)
         self.iter_timer.step(cur_time)
 
     def log_dict(self) -> dict[str, dict[str, int | float]]:

{xax-0.1.1 → xax-0.1.2/xax.egg-info}/PKG-INFO

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