xax 0.2.3__tar.gz → 0.2.5__tar.gz

{xax-0.2.3/xax.egg-info → xax-0.2.5}/PKG-INFO

@@ -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 → xax-0.2.5}/setup.py

@@ -15,14 +15,11 @@ with open("xax/requirements-dev.txt", "r", encoding="utf-8") as f:
     requirements_dev: list[str] = f.read().splitlines()
 
 requirements_export: list[str] = [
+    "flax",
     "orbax-export",
     "tensorflow",
 ]
 
-requirements_flax: list[str] = [
-    "flax",
-]
-
 with open("xax/__init__.py", "r", encoding="utf-8") as fh:
     version_re = re.search(r"^__version__ = \"([^\"]*)\"", fh.read(), re.MULTILINE)
 assert version_re is not None, "Could not find version in xax/__init__.py"
@@ -42,9 +39,8 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
-        "export": requirements_export,
-        "flax": requirements_flax,
-        "all": requirements_dev + requirements_export + requirements_flax,
+        "exportable": requirements_export,
+        "all": requirements_dev + requirements_export,
     },
     package_data={
         "xax": [

{xax-0.2.3 → xax-0.2.5}/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-0.2.3 → xax-0.2.5}/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-0.2.3 → xax-0.2.5}/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-0.2.3 → xax-0.2.5}/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-0.2.3 → xax-0.2.5}/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 → xax-0.2.5/xax.egg-info}/PKG-INFO

@@ -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 → xax-0.2.5}/xax.egg-info/requires.txt

@@ -23,9 +23,9 @@ pytest
 types-pillow
 types-psutil
 types-requests
+flax
 orbax-export
 tensorflow
-flax
 
 [dev]
 black
@@ -37,9 +37,7 @@ types-pillow
 types-psutil
 types-requests
 
-[export]
+[exportable]
+flax
 orbax-export
 tensorflow
-
-[flax]
-flax