xax 0.2.5__py3-none-any.whl → 0.2.6__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.5"
+__version__ = "0.2.6"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

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/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.dist-info → xax-0.2.6.dist-info}/METADATA

@@ -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.dist-info → xax-0.2.6.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=X_QqDNJir1wdsfRY1CU1F4mdCQMlMZnyqPtY8MM1ODU,14225
+xax/__init__.py,sha256=pvOmycOUli9h53XMTlinjJdY4zzw_fvWoF05UQwoItI,14225
 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
@@ -17,7 +17,7 @@ xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/base.py,sha256=OnXi2hiKPGwt6ng1dutnoQSiw7lEiWFlC_vx99_JsbQ,7694
-xax/task/logger.py,sha256=y4PGfMqKbfvPk8WCzr9MOsgG2X9E61KgeBVOYp-9kOY,40875
+xax/task/logger.py,sha256=gE67AaPCfU_1FpxY3t0yNRrIVgtp8Sax9UyOqFYMtzM,40976
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
 xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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=XcetJ0MppV_RDhgg1M9_d9heEXo-zeN_FS3MyczeBBU,31219
+xax/task/mixins/train.py,sha256=X3DIMPqpMRAJvFeeDaw5qFkH70R6yJelBAD8KHHwkUc,31196
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=d2H63ECtVOKySMUMrQRqq4kcuZpoXqo-L931usDVAhE,29903
@@ -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=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.5.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.5.dist-info/METADATA,sha256=4RBxZF_P0cg-a6QUNS9urvzc4BGGfoedqMrnP0L6Ksk,1879
-xax-0.2.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.5.dist-info/RECORD,,
+xax-0.2.6.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.6.dist-info/METADATA,sha256=g--2dJ3PIYDOWURjNhVmkv5B_VUcI2xY0iBa4oGuuVc,1879
+xax-0.2.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.2.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.6.dist-info/RECORD,,