xax 0.2.5__tar.gz → 0.2.6__tar.gz

{xax-0.2.5/xax.egg-info → xax-0.2.6}/PKG-INFO

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

{xax-0.2.5 → xax-0.2.6}/xax/__init__.py

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

{xax-0.2.5 → xax-0.2.6}/xax/task/logger.py

@@ -521,7 +521,8 @@ class LoggerImpl(ABC):
         Returns:
             If the logger should log the current step.
         """
-        return self.tickers[state.phase].tick(state.elapsed_time_s.item())
+        elapsed_time = state.elapsed_time_s.item() if state.phase == "train" else state.valid_elapsed_time_s.item()
+        return self.tickers[state.phase].tick(elapsed_time)
 
 
 class ToastHandler(logging.Handler):

{xax-0.2.5 → xax-0.2.6}/xax/task/mixins/train.py

@@ -115,19 +115,22 @@ class ValidStepTimer:
         self.last_valid_time: float | None = None
         self.last_valid_step: int | None = None
 
+    def _reset(self, state: State) -> None:
+        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:
         if state.num_steps < self.valid_first_n_steps:
             return True
 
         if self.last_valid_time is None or self.last_valid_step is None:
-            self.last_valid_time = state.elapsed_time_s.item()
-            self.last_valid_step = state.num_steps.item()
+            self._reset(state)
             return False
 
         # 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:
-            self.last_valid_step = state.num_steps.item()
+            self._reset(state)
             return True
 
         # Time-based validation.
@@ -136,14 +139,14 @@ class ValidStepTimer:
             valid_every_n_seconds is not None
             and state.elapsed_time_s.item() - self.last_valid_time >= valid_every_n_seconds
         ):
-            self.last_valid_time = state.elapsed_time_s.item()
+            self._reset(state)
             return True
 
         # 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:
-                self.last_valid_time = state.elapsed_time_s.item()
+                self._reset(state)
                 self.first_valid_step_flag = False
                 return True
 

{xax-0.2.5 → xax-0.2.6/xax.egg-info}/PKG-INFO

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