xax 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

xax/nn/embeddings.py

@@ -33,10 +33,10 @@ class LearnedPositionalEmbeddings(eqx.Module):
         learnable: Whether the embeddings are learnable.
     """
 
-    max_tsz: int = eqx.field(static=True)
-    embed_dim: int = eqx.field(static=True)
-    learnable: bool = eqx.field(static=True)
-    embeddings_tc: Array
+    max_tsz: int = eqx.field()
+    embed_dim: int = eqx.field()
+    learnable: bool = eqx.field()
+    embeddings_tc: Array = eqx.field()
 
     def __init__(
         self,
@@ -74,10 +74,10 @@ class SinusoidalEmbeddings(eqx.Module):
         base: The base for the sinusoidal embeddings.
     """
 
-    base: int = eqx.field(static=True)
-    max_tsz: int | None = eqx.field(static=True)
-    embed_dim: int | None = eqx.field(static=True)
-    embeddings_tc: Array | None
+    base: int = eqx.field()
+    max_tsz: int | None = eqx.field()
+    embed_dim: int | None = eqx.field()
+    embeddings_tc: Array | None = eqx.field()
 
     def __init__(
         self,
@@ -91,8 +91,8 @@ class SinusoidalEmbeddings(eqx.Module):
         self.max_tsz = max_tsz
         self.embed_dim = embed_dim
         self.base = base
+        self.embeddings_tc = None
 
-        self.embeddings_tc: Array | None = None
         if learnable:
             assert max_tsz is not None, "Learnable parameters require `max_tsz` to be set"
             assert embed_dim is not None, "Learnable parameters require `embed_dim` to be set"
@@ -192,7 +192,7 @@ class RotaryEmbeddings(eqx.Module):
         base: The base for the sinusoidal embeddings.
     """
 
-    base: int = eqx.field(static=True)
+    base: int = eqx.field()
 
     def __init__(self, base: int = 10_000) -> None:
         """Defines a rotary embeddings module.

xax/nn/geom.py

@@ -207,7 +207,7 @@ def quat_to_rotmat(quat: Array, eps: float = 1e-6) -> Array:
 
 def normalize(v: jnp.ndarray, axis: int = -1, eps: float = 1e-8) -> jnp.ndarray:
     norm = jnp.linalg.norm(v, axis=axis, keepdims=True)
-    return v / jnp.clip(norm, a_min=eps)
+    return v / jnp.clip(norm, min=eps)
 
 
 def rotation6d_to_rotation_matrix(r6d: jnp.ndarray) -> jnp.ndarray:
@@ -299,28 +299,28 @@ def rotation_matrix_to_quat(rotation_matrix: Array, eps: float = 1e-6) -> Array:
     trace = m00 + m11 + m22
 
     # Case 0: trace is positive
-    s0 = jnp.sqrt(jnp.clip(trace + 1.0, a_min=0.0)) * 2.0  # S = 4 * qw
+    s0 = jnp.sqrt(jnp.clip(trace + 1.0, min=0.0)) * 2.0  # S = 4 * qw
     w0 = 0.25 * s0
     x0 = (m21 - m12) / jnp.where(s0 < eps, 1.0, s0)
     y0 = (m02 - m20) / jnp.where(s0 < eps, 1.0, s0)
     z0 = (m10 - m01) / jnp.where(s0 < eps, 1.0, s0)
 
     # Case 1: m00 is the largest diagonal term
-    s1 = jnp.sqrt(jnp.clip(1.0 + m00 - m11 - m22, a_min=0.0)) * 2.0  # S = 4 * qx
+    s1 = jnp.sqrt(jnp.clip(1.0 + m00 - m11 - m22, min=0.0)) * 2.0  # S = 4 * qx
     w1 = (m21 - m12) / jnp.where(s1 < eps, 1.0, s1)
     x1 = 0.25 * s1
     y1 = (m01 + m10) / jnp.where(s1 < eps, 1.0, s1)
     z1 = (m02 + m20) / jnp.where(s1 < eps, 1.0, s1)
 
     # Case 2: m11 is the largest diagonal term
-    s2 = jnp.sqrt(jnp.clip(1.0 + m11 - m00 - m22, a_min=0.0)) * 2.0  # S = 4 * qy
+    s2 = jnp.sqrt(jnp.clip(1.0 + m11 - m00 - m22, min=0.0)) * 2.0  # S = 4 * qy
     w2 = (m02 - m20) / jnp.where(s2 < eps, 1.0, s2)
     x2 = (m01 + m10) / jnp.where(s2 < eps, 1.0, s2)
     y2 = 0.25 * s2
     z2 = (m12 + m21) / jnp.where(s2 < eps, 1.0, s2)
 
     # Case 3: m22 is the largest diagonal term
-    s3 = jnp.sqrt(jnp.clip(1.0 + m22 - m00 - m11, a_min=0.0)) * 2.0  # S = 4 * qz
+    s3 = jnp.sqrt(jnp.clip(1.0 + m22 - m00 - m11, min=0.0)) * 2.0  # S = 4 * qz
     w3 = (m10 - m01) / jnp.where(s3 < eps, 1.0, s3)
     x3 = (m02 + m20) / jnp.where(s3 < eps, 1.0, s3)
     y3 = (m12 + m21) / jnp.where(s3 < eps, 1.0, s3)

xax/nn/ssm.py

@@ -222,12 +222,12 @@ class DiscreteDiagSSMBlock(DiagSSMBlock):
 
 
 class SSM(eqx.Module):
-    vocab_embedding: eqx.nn.Embedding
-    output_layer: eqx.nn.Linear
-    blocks: list[BaseSSMBlock]
-    num_layers: int = eqx.static_field()
-    hidden_size: int = eqx.static_field()
-    skip_connections: bool = eqx.static_field()
+    vocab_embedding: eqx.nn.Embedding = eqx.field()
+    output_layer: eqx.nn.Linear = eqx.field()
+    blocks: list[BaseSSMBlock] = eqx.field()
+    num_layers: int = eqx.field()
+    hidden_size: int = eqx.field()
+    skip_connections: bool = eqx.field()
 
     def __init__(
         self,

xax/task/mixins/train.py

@@ -40,7 +40,12 @@ from xax.core.state import Phase, State
 from xax.nn.functions import set_random_seed
 from xax.nn.parallel import is_master
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
-from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin, CheckpointPart, load_ckpt
+from xax.task.mixins.checkpointing import (
+    CheckpointingConfig,
+    CheckpointingMixin,
+    CheckpointPart,
+    load_ckpt,
+)
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin

{xax-0.3.3.dist-info → xax-0.3.5.dist-info}/METADATA

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

{xax-0.3.3.dist-info → xax-0.3.5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=ffVd9_qSVuEAIPn6eK_6N8qEfDALRJigArHZQGy6y1o,15819
+xax/__init__.py,sha256=OeW6UObyosw6eJSEQ96AfRJIKHg5WyZ6xuZLJdcR6cg,16240
 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
@@ -8,14 +8,14 @@ xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
 xax/core/state.py,sha256=_gtINsRc310Bu_HuIYsDoOKTZa6DgU2tz0IOKkdnY9Q,3813
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/nn/attention.py,sha256=0essK90OO3x9FxnUqU0DhufwXKRMN41zMtRCki5iKzQ,24742
-xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
+xax/nn/attention.py,sha256=ESO6THJ5ORKxSM8LziRLEkj1d_QXtDndPi80Puyo-xA,28033
+xax/nn/embeddings.py,sha256=8tAuAPdkVj-U5IwtRZKHA0WYMFRbpCuwyAxcChdKhbE,11784
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
-xax/nn/geom.py,sha256=6rBQrZRX1miG08VG-s8phPjA6MEFxUAfQVPt5F0RQQI,10645
+xax/nn/geom.py,sha256=c9K52vLm-V-15CRqMNx0OmqsWfb3PHQxXW4OSx9kCAk,10635
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/metrics.py,sha256=zuvPXlRQczBTLHD4ilNGmZaiq6Yie3rxCMq6JkI_kos,3154
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
-xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
+xax/nn/ssm.py,sha256=qSBv_FobnaFA5jt87OF5P2q5ih6sj4SlehhEhEFaPjA,10766
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/base.py,sha256=i6FRJ75aqlekWkzJNRWDUEX7P514pUjLVuxjhX1GBgw,8198
 xax/task/logger.py,sha256=Bmhl4mv08Aq49ZyX6BdjPIsPJK28e8s3mVFatM4IY2Q,41060
@@ -42,7 +42,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=pcLrYc_TycZUY9zZim05Skc2FWk3IZKFnu6p3UDMonM,1966
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=TZatz5QwTfrNhQTiO2IqrmQY9P4Lay6FAD2VsQpWa54,33245
+xax/task/mixins/train.py,sha256=bjBoigTCjbq9H4hcqIO32irHBc9rC2zkgXrnGNI2RtI,33266
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=5k5hPYSaVjzoR_nm2Q3DAHMMYi3Bcp3N3PAQbwZq7Gg,29830
@@ -59,9 +59,9 @@ 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.3.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.3.dist-info/METADATA,sha256=mjIzoFZDSR3V1-2LHbvup6wDVa4vLiqbqiNWLsKCXY8,1246
-xax-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.3.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.3.dist-info/RECORD,,
+xax-0.3.5.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.5.dist-info/METADATA,sha256=kMRKGih6o7SfqGrvGQW_7OkFST6PDnbPuopnfx_bAOs,1246
+xax-0.3.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.5.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.5.dist-info/RECORD,,