xax 0.3.2__tar.gz → 0.3.4__tar.gz

{xax-0.3.2/xax.egg-info → xax-0.3.4}/PKG-INFO

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

{xax-0.3.2 → xax-0.3.4}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.2"
+__version__ = "0.3.4"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,10 +23,14 @@ __all__ = [
     "get_run_dir",
     "load_user_config",
     "State",
+    "AttentionCache",
+    "AttentionCacheDict",
     "CrossAttentionBlock",
     "SelfAttentionBlock",
     "Transformer",
     "TransformerBlock",
+    "TransformerCache",
+    "TransformerStack",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -42,11 +46,12 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "normalize",
+    "quat_mul",
     "quat_to_euler",
     "quat_to_rotmat",
-    "quat_mul",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
+    "rotation_matrix_to_quat",
     "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
@@ -205,10 +210,14 @@ NAME_MAP: dict[str, str] = {
     "get_run_dir": "core.conf",
     "load_user_config": "core.conf",
     "State": "core.state",
+    "AttentionCache": "nn.attention",
+    "AttentionCacheDict": "nn.attention",
     "CrossAttentionBlock": "nn.attention",
     "SelfAttentionBlock": "nn.attention",
     "Transformer": "nn.attention",
     "TransformerBlock": "nn.attention",
+    "TransformerCache": "nn.attention",
+    "TransformerStack": "nn.attention",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -224,11 +233,12 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "normalize": "nn.geom",
+    "quat_mul": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
-    "quat_mul": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
     "rotation6d_to_rotation_matrix": "nn.geom",
+    "rotation_matrix_to_quat": "nn.geom",
     "rotation_matrix_to_rotation6d": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.metrics",
@@ -360,6 +370,9 @@ NAME_MAP.update(
     },
 )
 
+# In NAME_MAP
+NAME_MAP["TransformerStack"] = "nn.attention"
+
 
 def __getattr__(name: str) -> object:
     if name not in NAME_MAP:
@@ -380,7 +393,16 @@ 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.attention import (
+        AttentionCache,
+        AttentionCacheDict,
+        CrossAttentionBlock,
+        SelfAttentionBlock,
+        Transformer,
+        TransformerBlock,
+        TransformerCache,
+        TransformerStack,
+    )
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -405,15 +427,11 @@ if IMPORT_ALL or TYPE_CHECKING:
         quat_to_rotmat,
         rotate_vector_by_quat,
         rotation6d_to_rotation_matrix,
+        rotation_matrix_to_quat,
         rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy
-    from xax.nn.metrics import (
-        NormType,
-        cast_norm_type,
-        dynamic_time_warping,
-        get_norm,
-    )
+    from xax.nn.metrics import NormType, cast_norm_type, dynamic_time_warping, get_norm
     from xax.nn.parallel import is_master
     from xax.nn.ssm import SSM, BaseSSMBlock, DiagSSMBlock, SSMBlock
     from xax.task.base import RawConfigType