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

xax/core/state.py

@@ -1,6 +1,5 @@
 """Defines a dataclass for keeping track of the current training state."""
 
-import time
 from dataclasses import dataclass
 from typing import Literal, NotRequired, TypedDict, Unpack, cast
 
@@ -19,6 +18,8 @@ def _phase_to_int(phase: Phase) -> int:
 
 
 def _int_to_phase(i: int) -> Phase:
+    if i < 0 or i > 1:
+        raise ValueError(f"Invalid phase: {i}")
     return cast(Phase, ["train", "valid"][i])
 
 
@@ -27,8 +28,8 @@ class StateDict(TypedDict, total=False):
     num_samples: NotRequired[int | Array]
     num_valid_steps: NotRequired[int | Array]
     num_valid_samples: NotRequired[int | Array]
-    start_time_s: NotRequired[float | Array]
     elapsed_time_s: NotRequired[float | Array]
+    valid_elapsed_time_s: NotRequired[float | Array]
     phase: NotRequired[Phase]
     _phase: NotRequired[int | Array]
 
@@ -43,24 +44,24 @@ class State:
     def num_steps(self) -> Array:
         return self._int32_arr[0]
 
-    @property
-    def num_samples(self) -> Array:
-        return self._float32_arr[0]
-
     @property
     def num_valid_steps(self) -> Array:
         return self._int32_arr[1]
 
+    @property
+    def num_samples(self) -> Array:
+        return self._float32_arr[0]
+
     @property
     def num_valid_samples(self) -> Array:
         return self._float32_arr[1]
 
     @property
-    def start_time_s(self) -> Array:
+    def elapsed_time_s(self) -> Array:
         return self._float32_arr[2]
 
     @property
-    def elapsed_time_s(self) -> Array:
+    def valid_elapsed_time_s(self) -> Array:
         return self._float32_arr[3]
 
     @property
@@ -71,7 +72,7 @@ class State:
     def init_state(cls) -> "State":
         return cls(
             _int32_arr=jnp.array([0, 0, 0], dtype=jnp.int32),
-            _float32_arr=jnp.array([0.0, 0.0, time.time(), 0.0], dtype=jnp.float32),
+            _float32_arr=jnp.array([0.0, 0.0, 0.0, 0.0], dtype=jnp.float32),
         )
 
     @property
@@ -88,19 +89,19 @@ class State:
             int32_arr = int32_arr.at[1].set(kwargs["num_valid_steps"])
 
         if "phase" in kwargs:
-            int32_arr = int32_arr.at[3].set(_phase_to_int(kwargs["phase"]))
+            int32_arr = int32_arr.at[2].set(_phase_to_int(kwargs["phase"]))
         if "_phase" in kwargs:
-            int32_arr = int32_arr.at[3].set(kwargs["_phase"])
+            int32_arr = int32_arr.at[2].set(kwargs["_phase"])
 
         if "num_samples" in kwargs:
             float32_arr = float32_arr.at[0].set(kwargs["num_samples"])
         if "num_valid_samples" in kwargs:
             float32_arr = float32_arr.at[1].set(kwargs["num_valid_samples"])
 
-        if "start_time_s" in kwargs:
-            float32_arr = float32_arr.at[2].set(kwargs["start_time_s"])
         if "elapsed_time_s" in kwargs:
-            float32_arr = float32_arr.at[3].set(kwargs["elapsed_time_s"])
+            float32_arr = float32_arr.at[2].set(kwargs["elapsed_time_s"])
+        if "valid_elapsed_time_s" in kwargs:
+            float32_arr = float32_arr.at[3].set(kwargs["valid_elapsed_time_s"])
 
         return State(
             _int32_arr=int32_arr,
@@ -110,11 +111,11 @@ class State:
     def to_dict(self) -> dict[str, int | float | str]:
         return {
             "num_steps": int(self.num_steps),
-            "num_samples": int(self.num_samples),
             "num_valid_steps": int(self.num_valid_steps),
+            "num_samples": int(self.num_samples),
             "num_valid_samples": int(self.num_valid_samples),
-            "start_time_s": float(self.start_time_s),
             "elapsed_time_s": float(self.elapsed_time_s),
+            "valid_elapsed_time_s": float(self.valid_elapsed_time_s),
             "phase": str(self.phase),
         }
 
@@ -126,9 +127,7 @@ class State:
         int32_arr = jnp.array(
             [
                 d.get("num_steps", 0),
-                d.get("num_samples", 0),
                 d.get("num_valid_steps", 0),
-                d.get("num_valid_samples", 0),
                 d.get("_phase", 0),
             ],
             dtype=jnp.int32,
@@ -136,8 +135,10 @@ class State:
 
         float32_arr = jnp.array(
             [
-                d.get("start_time_s", time.time()),
+                d.get("num_samples", 0),
+                d.get("num_valid_samples", 0),
                 d.get("elapsed_time_s", 0.0),
+                d.get("valid_elapsed_time_s", 0.0),
             ],
             dtype=jnp.float32,
         )

xax/nn/export.py

@@ -14,8 +14,8 @@ try:
     from orbax.export import ExportManager, JaxModule, ServingConfig
 except ImportError as e:
     raise ImportError(
-        "In order to export models, please install Xax with export dependencies, "
-        "using 'xax[export]` to install the required dependencies."
+        "In order to export models, please install Xax with exportable dependencies, "
+        "using 'xax[exportable]` to install the required dependencies."
     ) from e
 
 logger = logging.getLogger(__name__)

xax/task/loggers/tensorboard.py

@@ -157,7 +157,7 @@ class TensorboardLogger(LoggerImpl):
         writer = self.get_writer(line.state.phase)
 
         global_step = line.state.num_steps.item()
-        walltime = (line.state.start_time_s + line.state.elapsed_time_s).item()
+        walltime = line.state.elapsed_time_s.item()
 
         for namespace, scalars in line.scalars.items():
             for scalar_key, scalar_value in scalars.items():

xax/task/mixins/train.py

@@ -720,20 +720,21 @@ class TrainMixin(
 
         while not self.is_training_over(state):
             if self.valid_step_timer.is_valid_step(state):
-                valid_batch = next(valid_pf)
+                with ContextTimer() as timer:
+                    valid_batch = next(valid_pf)
+                    output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
+                    self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
+
                 state = state.replace(
                     phase="valid",
                     num_valid_steps=state.num_valid_steps + 1,
                     num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
+                    valid_elapsed_time_s=state.valid_elapsed_time_s + timer.elapsed_time,
                 )
 
-                output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
-                self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
-
-            state = self.on_step_start(state)
-            train_batch = next(train_pf)
-
             with ContextTimer() as timer:
+                state = self.on_step_start(state)
+                train_batch = next(train_pf)
                 model_arr, opt_state, output, metrics = self.train_step(
                     model_arr=model_arr,
                     model_static=model_static,

{xax-0.2.3.dist-info → xax-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.3
+Version: 0.2.5
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,11 +31,10 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Provides-Extra: export
-Requires-Dist: orbax-export; extra == "export"
-Requires-Dist: tensorflow; extra == "export"
-Provides-Extra: flax
-Requires-Dist: flax; extra == "flax"
+Provides-Extra: exportable
+Requires-Dist: flax; extra == "exportable"
+Requires-Dist: orbax-export; extra == "exportable"
+Requires-Dist: tensorflow; extra == "exportable"
 Provides-Extra: all
 Requires-Dist: black; extra == "all"
 Requires-Dist: darglint; extra == "all"
@@ -45,9 +44,9 @@ Requires-Dist: pytest; extra == "all"
 Requires-Dist: types-pillow; extra == "all"
 Requires-Dist: types-psutil; extra == "all"
 Requires-Dist: types-requests; extra == "all"
+Requires-Dist: flax; extra == "all"
 Requires-Dist: orbax-export; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: flax; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.3.dist-info → xax-0.2.5.dist-info}/RECORD

@@ -1,14 +1,14 @@
-xax/__init__.py,sha256=P4q2IGkfpHaN3ZlGFiW0bzWm1spLSUyl0GEPvH8oITg,14225
+xax/__init__.py,sha256=X_QqDNJir1wdsfRY1CU1F4mdCQMlMZnyqPtY8MM1ODU,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
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=lA7A5HCm2Nwk4J0kJGlTIhqHYFvbuwHfdNzOhmjEW08,4453
+xax/core/state.py,sha256=yO25lMoLCUTJlHyLzQxlDbsHC_GZ3HkrKAq5huA7AkU,4552
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
-xax/nn/export.py,sha256=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
+xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
 xax/nn/geom.py,sha256=rImNlkHWeoNcY7f84nknizJ6uzsrMhbAtKeb2xAWxNY,6215
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
@@ -29,7 +29,7 @@ xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,
 xax/task/loggers/json.py,sha256=_tKum6jk_gqVzO-4MqSNXbE-Mmn-yJzkRAT-N1y2zes,4139
 xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
 xax/task/loggers/stdout.py,sha256=ERLFrYe61hSSztzyxBRseobHQR72YFYjEd2i_hOeJ20,6595
-xax/task/loggers/tensorboard.py,sha256=3ohI6STgSCbU8oyeiH_f3QyLVF_zO_6dwjn0ns59rUU,8334
+xax/task/loggers/tensorboard.py,sha256=KFlsK0zD2ubDqAXYL4Ds7NQ9F-Ke-PHwfhLOYsGcbw4,8306
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
 xax/task/mixins/checkpointing.py,sha256=8Hi-2G0EA5OFRjgiOutlk7HgkD5b-0GHazOAYxnGytM,11409
@@ -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=lMHCnxsbZJbwK3esL5S3cJ0Jf5Qx19Y4pm3A7NY-TIE,31064
+xax/task/mixins/train.py,sha256=XcetJ0MppV_RDhgg1M9_d9heEXo-zeN_FS3MyczeBBU,31219
 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.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.3.dist-info/METADATA,sha256=ukAnG444wnzRpXgmHSrs7RKJ-UQvOdl6ZE2ZrN0w4Yg,1882
-xax-0.2.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.3.dist-info/RECORD,,
+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,,