xax 0.2.10__tar.gz → 0.2.11__tar.gz

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

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

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

{xax-0.2.10 → xax-0.2.11}/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-0.2.10 → xax-0.2.11}/xax/task/mixins/train.py

@@ -121,7 +121,7 @@ class ValidStepTimer:
         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:
+        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:
@@ -272,7 +272,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,

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

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