xax 0.3.3__tar.gz → 0.3.9__tar.gz

{xax-0.3.3/xax.egg-info → xax-0.3.9}/PKG-INFO

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

{xax-0.3.3 → xax-0.3.9}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.3"
+__version__ = "0.3.9"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,10 +23,18 @@ __all__ = [
     "get_run_dir",
     "load_user_config",
     "State",
+    "AttentionCache",
+    "AttentionCacheDict",
     "CrossAttentionBlock",
     "SelfAttentionBlock",
     "Transformer",
     "TransformerBlock",
+    "TransformerCache",
+    "TransformerStack",
+    "Categorical",
+    "Distribution",
+    "MixtureOfGaussians",
+    "Normal",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -132,6 +140,7 @@ __all__ = [
     "compute_nan_ratio",
     "flatten_array",
     "flatten_pytree",
+    "get_pytree_mapping",
     "get_pytree_param_count",
     "pytree_has_nans",
     "reshuffle_pytree",
@@ -206,10 +215,18 @@ 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",
+    "Categorical": "nn.distributions",
+    "Distribution": "nn.distributions",
+    "MixtureOfGaussians": "nn.distributions",
+    "Normal": "nn.distributions",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -315,6 +332,7 @@ NAME_MAP: dict[str, str] = {
     "compute_nan_ratio": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
+    "get_pytree_mapping": "utils.pytree",
     "get_pytree_param_count": "utils.pytree",
     "pytree_has_nans": "utils.pytree",
     "reshuffle_pytree": "utils.pytree",
@@ -362,6 +380,9 @@ NAME_MAP.update(
     },
 )
 
+# In NAME_MAP
+NAME_MAP["TransformerStack"] = "nn.attention"
+
 
 def __getattr__(name: str) -> object:
     if name not in NAME_MAP:
@@ -382,7 +403,17 @@ 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.distributions import Categorical, Distribution, MixtureOfGaussians, Normal
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -411,12 +442,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         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
@@ -494,6 +520,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         compute_nan_ratio,
         flatten_array,
         flatten_pytree,
+        get_pytree_mapping,
         get_pytree_param_count,
         pytree_has_nans,
         reshuffle_pytree,