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

xax/task/loggers/tensorboard.py

@@ -32,7 +32,7 @@ class TensorboardLogger(LoggerImpl):
         flush_seconds: float = 10.0,
         wait_seconds: float = 0.0,
         start_in_subprocess: bool = True,
-        use_localhost: bool = False,
+        use_localhost: bool | None = None,
         log_interval_seconds: float = 10.0,
     ) -> None:
         """Defines a logger which writes to Tensorboard.
@@ -49,6 +49,9 @@ class TensorboardLogger(LoggerImpl):
         """
         super().__init__(log_interval_seconds)
 
+        if use_localhost is None:
+            use_localhost = os.environ.get("USE_TENSORBOARD_LOCALHOST", "0") == "1"
+
         self.log_directory = Path(run_directory).expanduser().resolve() / subdirectory
         self.wait_seconds = wait_seconds
         self.start_in_subprocess = start_in_subprocess

xax/task/mixins/train.py

@@ -120,8 +120,8 @@ 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:
-        if state.num_steps < self.valid_first_n_steps:
+    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
 
         if self.last_valid_time is None or self.last_valid_step is None:
@@ -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
@@ -272,7 +278,7 @@ class TrainMixin(
                 v, secondary = v
             else:
                 secondary = False
-            self.logger.log_scalar(k, v, namespace="⌛ timers", secondary=secondary)
+            self.logger.log_scalar(k, v, namespace="🕒 timers", secondary=secondary)
 
     def log_step(
         self,
@@ -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.10.dist-info → xax-0.2.12.dist-info}/METADATA

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

{xax-0.2.10.dist-info → xax-0.2.12.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=8gVPgaF2I-ObP5RU58OEvxPtHOD5a1i382CZbuW9jvM,15510
+xax/__init__.py,sha256=GprJJQKvQtJp16V1-8GYWLgUtwueKLRsOrtWKUZRq9o,15510
 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
@@ -29,7 +29,7 @@ xax/task/loggers/callback.py,sha256=zQuV1xCvz47Q3UQqP1D5mBhbVzptvmPR_7hX25vqSk0,
 xax/task/loggers/json.py,sha256=6A5wL7kspsXnpPhI_vu0scgd2Z2-WLhw4gbBFm7eZMM,4377
 xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,1407
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
-xax/task/loggers/tensorboard.py,sha256=gkAalLsYPGjZiiMlqvDWIhNpYCfKWNvnPz3brIv3JaQ,8725
+xax/task/loggers/tensorboard.py,sha256=sdsA8GjZG5JQpoAxNDRr_bGvqN8Olgj_almZBb2K5F8,8850
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
 xax/task/mixins/checkpointing.py,sha256=zqospBFnTbGt_iriiduVfXazINPbzWpwmIs91KAniMY,10147
@@ -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=ln-Qvx27gruRCyaVLHvTL1XzCDoEJ53Ep9Gp3JwsT9s,31268
+xax/task/mixins/train.py,sha256=F5hNHGXCFQ2QbUqdZqhWcpDzR1JVoBOo3oCm0tcpAlA,31605
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=bj8BftSHT3fFzfiJ0Co0WvqWo0rUS8kQnQYpVvH8FTM,29942
@@ -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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.10.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.10.dist-info/METADATA,sha256=PwQsQ3uDewGQvQ5XQ6NT4SgZ15woG6YZgMw1ejLRT4g,1880
-xax-0.2.10.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-xax-0.2.10.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.10.dist-info/RECORD,,
+xax-0.2.12.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.12.dist-info/METADATA,sha256=VJjr__ccXjzqrpvDzevgRlDW22GwLQGqAfTsHmbEk7I,1880
+xax-0.2.12.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+xax-0.2.12.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.12.dist-info/RECORD,,