xax 0.2.11__tar.gz → 0.2.12__tar.gz

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

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

{xax-0.2.11 → xax-0.2.12}/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.12"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.2.11 → xax-0.2.12}/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,7 +728,7 @@ 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):
+            if self.valid_step_timer(state):
                 with ContextTimer() as timer:
                     valid_batch = next(valid_pf)
                     output, metrics = self.val_step(model_arr, model_static, valid_batch, state)

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

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