xax 0.2.21__tar.gz → 0.2.23__tar.gz

{xax-0.2.21/xax.egg-info → xax-0.2.23}/PKG-INFO

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

{xax-0.2.21 → xax-0.2.23}/pyproject.toml

@@ -38,6 +38,7 @@ module = [
     "optax.*",
     "setuptools.*",
     "tensorboard.*",
+    "tensorflow_datasets.*",
     "transformers.*",
 ]
 

{xax-0.2.21 → xax-0.2.23}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.21"
+__version__ = "0.2.23"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,6 +23,10 @@ __all__ = [
     "get_run_dir",
     "load_user_config",
     "State",
+    "CrossAttentionBlock",
+    "SelfAttentionBlock",
+    "Transformer",
+    "TransformerBlock",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -112,8 +116,10 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
+    "grad",
     "jit",
     "scan",
+    "vmap",
     "save_jaxpr_dot",
     "ColoredFormatter",
     "configure_logging",
@@ -198,6 +204,10 @@ NAME_MAP: dict[str, str] = {
     "get_run_dir": "core.conf",
     "load_user_config": "core.conf",
     "State": "core.state",
+    "CrossAttentionBlock": "nn.attention",
+    "SelfAttentionBlock": "nn.attention",
+    "Transformer": "nn.attention",
+    "TransformerBlock": "nn.attention",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -287,8 +297,10 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
+    "grad": "utils.jax",
     "jit": "utils.jax",
     "scan": "utils.jax",
+    "vmap": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
@@ -366,6 +378,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         load_user_config,
     )
     from xax.core.state import Phase, State
+    from xax.nn.attention import CrossAttentionBlock, SelfAttentionBlock, Transformer, TransformerBlock
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -460,7 +473,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import jit, scan
+    from xax.utils.jax import grad, jit, scan, vmap
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,

{xax-0.2.21 → xax-0.2.23}/xax/core/state.py

@@ -27,11 +27,8 @@ def _int_to_phase(i: int) -> Phase:
 class StateDict(TypedDict, total=False):
     num_steps: NotRequired[int | Array]
     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]
 
@@ -47,38 +44,26 @@ class State:
         return self._int32_arr[0]
 
     @property
-    def num_valid_steps(self) -> Array:
-        return self._int32_arr[1]
+    def phase(self) -> Phase:
+        return _int_to_phase(self._int32_arr[1].item())
 
     @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:
-        return self._float32_arr[2]
+        return self._float32_arr[1]
 
     @property
     def elapsed_time_s(self) -> Array:
-        return self._float32_arr[3]
-
-    @property
-    def valid_elapsed_time_s(self) -> Array:
-        return self._float32_arr[4]
-
-    @property
-    def phase(self) -> Phase:
-        return _int_to_phase(self._int32_arr[2].item())
+        return self._float32_arr[2]
 
     @classmethod
     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, 0.0], dtype=jnp.float32),
+            _int32_arr=jnp.array([0, 0], dtype=jnp.int32),
+            _float32_arr=jnp.array([0.0, time.time(), 0.0], dtype=jnp.float32),
         )
 
     @property
@@ -91,25 +76,19 @@ class State:
 
         if "num_steps" in kwargs:
             int32_arr = int32_arr.at[0].set(kwargs["num_steps"])
-        if "num_valid_steps" in kwargs:
-            int32_arr = int32_arr.at[1].set(kwargs["num_valid_steps"])
 
         if "phase" in kwargs:
-            int32_arr = int32_arr.at[2].set(_phase_to_int(kwargs["phase"]))
+            int32_arr = int32_arr.at[1].set(_phase_to_int(kwargs["phase"]))
         if "_phase" in kwargs:
-            int32_arr = int32_arr.at[2].set(kwargs["_phase"])
+            int32_arr = int32_arr.at[1].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"])
+            float32_arr = float32_arr.at[1].set(kwargs["start_time_s"])
         if "elapsed_time_s" in kwargs:
-            float32_arr = float32_arr.at[3].set(kwargs["elapsed_time_s"])
-        if "valid_elapsed_time_s" in kwargs:
-            float32_arr = float32_arr.at[4].set(kwargs["valid_elapsed_time_s"])
+            float32_arr = float32_arr.at[2].set(kwargs["elapsed_time_s"])
 
         return State(
             _int32_arr=int32_arr,
@@ -119,12 +98,9 @@ class State:
     def to_dict(self) -> dict[str, int | float | str]:
         return {
             "num_steps": int(self.num_steps.item()),
-            "num_valid_steps": int(self.num_valid_steps.item()),
             "num_samples": int(self.num_samples.item()),
-            "num_valid_samples": int(self.num_valid_samples.item()),
             "start_time_s": float(self.start_time_s.item()),
             "elapsed_time_s": float(self.elapsed_time_s.item()),
-            "valid_elapsed_time_s": float(self.valid_elapsed_time_s.item()),
             "phase": str(self.phase),
         }
 
@@ -136,7 +112,6 @@ class State:
         int32_arr = jnp.array(
             [
                 d.get("num_steps", 0),
-                d.get("num_valid_steps", 0),
                 d.get("_phase", 0),
             ],
             dtype=jnp.int32,
@@ -145,10 +120,8 @@ class State:
         float32_arr = jnp.array(
             [
                 d.get("num_samples", 0),
-                d.get("num_valid_samples", 0),
                 d.get("start_time_s", time.time()),
                 d.get("elapsed_time_s", 0.0),
-                d.get("valid_elapsed_time_s", 0.0),
             ],
             dtype=jnp.float32,
         )