xax 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.6"
+__version__ = "0.1.7"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -41,9 +41,6 @@ __all__ = [
     "load_eqx_mlp",
     "make_eqx_mlp",
     "save_eqx",
-    "export",
-    "export_flax",
-    "export_with_params",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
@@ -97,8 +94,6 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
-    "HashableArray",
-    "hashable_array",
     "jit",
     "save_jaxpr_dot",
     "ColoredFormatter",
@@ -132,6 +127,9 @@ __all__ = [
     "snakecase_to_camelcase",
     "uncolored",
     "wrapped",
+    "FrozenDict",
+    "HashableArray",
+    "hashable_array",
 ]
 
 __all__ += [
@@ -161,7 +159,7 @@ if "XLA_FLAGS" in os.environ:
 # If Nvidia GPU is detected (meaning, is `nvidia-smi` available?), disable
 # Triton GEMM kernels. See https://github.com/NVIDIA/JAX-Toolbox
 if shutil.which("nvidia-smi") is not None:
-    xla_flags += ["--xla_gpu_enable_latency_hiding_scheduler", "--xla_gpu_enable_triton_gemm"]
+    xla_flags += ["--xla_gpu_enable_latency_hiding_scheduler=true", "--xla_gpu_enable_triton_gemm=false"]
 os.environ["XLA_FLAGS"] = " ".join(xla_flags)
 
 # If this flag is set, eagerly imports the entire package (not recommended).
@@ -197,9 +195,6 @@ NAME_MAP: dict[str, str] = {
     "load_eqx_mlp": "nn.equinox",
     "make_eqx_mlp": "nn.equinox",
     "save_eqx": "nn.equinox",
-    "export": "nn.export",
-    "export_flax": "nn.export",
-    "export_with_params": "nn.export",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
@@ -253,8 +248,6 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
-    "HashableArray": "utils.jax",
-    "hashable_array": "utils.jax",
     "jit": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
@@ -288,6 +281,9 @@ NAME_MAP: dict[str, str] = {
     "snakecase_to_camelcase": "utils.text",
     "uncolored": "utils.text",
     "wrapped": "utils.text",
+    "FrozenDict": "utils.types.frozen_dict",
+    "HashableArray": "utils.types.hashable_array",
+    "hashable_array": "utils.types.hashable_array",
 }
 
 # Need to manually set some values which can't be auto-generated.
@@ -352,11 +348,6 @@ if IMPORT_ALL or TYPE_CHECKING:
         make_eqx_mlp,
         save_eqx,
     )
-    from xax.nn.export import (
-        export,
-        export_flax,
-        export_with_params,
-    )
     from xax.nn.geom import (
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
@@ -408,7 +399,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import HashableArray, hashable_array, jit
+    from xax.utils.jax import jit
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,
@@ -448,5 +439,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         uncolored,
         wrapped,
     )
+    from xax.utils.types.frozen_dict import FrozenDict
+    from xax.utils.types.hashable_array import HashableArray, hashable_array
 
 del TYPE_CHECKING, IMPORT_ALL

xax/nn/equinox.py

@@ -72,7 +72,7 @@ def _infer_activation(activation: ActivationFunction) -> Callable:
         raise ValueError(f"Activation function `{activation}` not found in `jax.nn`")
 
 
-def make_eqx_mlp(hyperparams: MLPHyperParams, key: PRNGKeyArray = jax.random.PRNGKey(0)) -> eqx.nn.MLP:
+def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.MLP:
     """Create an Equinox MLP from a set of hyperparameters.
 
     Args:
@@ -176,5 +176,5 @@ def load_eqx_mlp(
 ) -> eqx.nn.MLP:
     with open(eqx_file, "rb") as f:
         hyperparams = json.loads(f.readline().decode(encoding="utf-8"))
-        model = make_eqx_mlp(hyperparams=hyperparams)
+        model = make_eqx_mlp(hyperparams=hyperparams, key=jax.random.PRNGKey(0))
         return eqx.tree_deserialise_leaves(f, model)

xax/nn/export.py

@@ -4,18 +4,18 @@ import logging
 from pathlib import Path
 from typing import Callable
 
-import flax
 import jax
-import tensorflow as tf
-from jax.experimental import jax2tf
 from jaxtyping import Array, PyTree
 
 try:
+    import flax
+    import tensorflow as tf
+    from jax.experimental import jax2tf
     from orbax.export import ExportManager, JaxModule, ServingConfig
 except ImportError as e:
     raise ImportError(
-        "Please install the package with `orbax` as a dependency, using "
-        "'xax[export]` to install the required dependencies."
+        "In order to export models, please install Xax with export dependencies, "
+        "using 'xax[export]` to install the required dependencies."
     ) from e
 
 logger = logging.getLogger(__name__)

xax/task/mixins/compile.py

@@ -22,7 +22,7 @@ def get_cache_dir() -> str | None:
     # By default, only cache on MacOS, since Jax caching on Linux is very
     # prone to NaNs.
     match sys.platform:
-        case "darwin":
+        case "darwin" | "linux":
             return str((Path.home() / ".cache" / "jax" / "jaxcache").resolve())
         case _:
             return None

xax/utils/jax.py

@@ -138,25 +138,3 @@ def jit(
         return wrapped
 
     return decorator
-
-
-class HashableArray:
-    def __init__(self, array: np.ndarray | jnp.ndarray) -> None:
-        if not isinstance(array, (np.ndarray, jnp.ndarray)):
-            raise ValueError(f"Expected np.ndarray or jnp.ndarray, got {type(array)}")
-        self.array = array
-        self._hash: int | None = None
-
-    def __hash__(self) -> int:
-        if self._hash is None:
-            self._hash = hash(self.array.tobytes())
-        return self._hash
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, HashableArray):
-            return False
-        return bool(jnp.array_equal(self.array, other.array))
-
-
-def hashable_array(array: np.ndarray | jnp.ndarray) -> HashableArray:
-    return HashableArray(array)

xax/utils/types/frozen_dict.py

@@ -0,0 +1,148 @@
+"""Defines a frozen dictionary type.
+
+This is mostly taken from Flax - we move it here to avoid having to use Flax as
+a dependency in downstream projects.
+"""
+
+import collections
+from types import MappingProxyType
+from typing import Any, Iterator, Mapping, Self, TypeVar
+
+import jax
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+def _prepare_freeze(xs: Any) -> Any:  # noqa: ANN401
+    """Deep copy unfrozen dicts to make the dictionary FrozenDict safe."""
+    if isinstance(xs, FrozenDict):
+        return xs._dict  # pylint: disable=protected-access
+    if not isinstance(xs, dict):
+        return xs
+    return {key: _prepare_freeze(val) for key, val in xs.items()}
+
+
+def _indent(x: str, num_spaces: int) -> str:
+    indent_str = " " * num_spaces
+    lines = x.split("\n")
+    assert not lines[-1]
+    return "\n".join(indent_str + line for line in lines[:-1]) + "\n"
+
+
+class FrozenKeysView(collections.abc.KeysView[K]):
+    def __repr__(self) -> str:
+        return f"frozen_dict_keys({list(self)})"
+
+
+class FrozenValuesView(collections.abc.ValuesView[V]):
+    def __repr__(self) -> str:
+        return f"frozen_dict_values({list(self)})"
+
+
+@jax.tree_util.register_pytree_with_keys_class
+class FrozenDict(Mapping[K, V]):
+    """An immutable variant of the Python dict."""
+
+    __slots__ = ("_dict", "_hash")
+
+    def __init__(self, *args: Any, __unsafe_skip_copy__: bool = False, **kwargs: Any) -> None:  # noqa: ANN401
+        # make sure the dict is as
+        xs = dict(*args, **kwargs)
+        if __unsafe_skip_copy__:
+            self._dict = xs
+        else:
+            self._dict = _prepare_freeze(xs)
+
+        self._hash: int | None = None
+
+    def __getitem__(self, key: K) -> V:
+        v = self._dict[key]
+        if isinstance(v, dict):
+            return FrozenDict(v)  # type: ignore[return-value]
+        return v
+
+    def __setitem__(self, key: K, value: V) -> None:
+        raise ValueError("FrozenDict is immutable.")
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._dict
+
+    def __iter__(self) -> Iterator[K]:
+        return iter(self._dict)
+
+    def __len__(self) -> int:
+        return len(self._dict)
+
+    def __repr__(self) -> str:
+        return self.pretty_repr()
+
+    def __reduce__(self) -> tuple[type["FrozenDict[K, V]"], tuple[dict[K, V]]]:
+        return FrozenDict, (self.unfreeze(),)
+
+    def pretty_repr(self, num_spaces: int = 4) -> str:
+        """Returns an indented representation of the nested dictionary."""
+
+        def pretty_dict(x: Any) -> str:  # noqa: ANN401
+            if not isinstance(x, dict):
+                return repr(x)
+            rep = ""
+            for key, val in x.items():
+                rep += f"{key}: {pretty_dict(val)},\n"
+            if rep:
+                return "{\n" + _indent(rep, num_spaces) + "}"
+            else:
+                return "{}"
+
+        return f"FrozenDict({pretty_dict(self._dict)})"
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            h = 0
+            for key, value in self.items():
+                h ^= hash((key, value))
+            self._hash = h
+        return self._hash
+
+    def copy(self, add_or_replace: Mapping[K, V] = MappingProxyType({})) -> Self:
+        return type(self)({**self, **unfreeze(add_or_replace)})  # type: ignore[arg-type]
+
+    def keys(self) -> FrozenKeysView[K]:
+        return FrozenKeysView(self)
+
+    def values(self) -> FrozenValuesView[V]:
+        return FrozenValuesView(self)
+
+    def items(self) -> Iterator[tuple[K, V]]:  # type: ignore[override]
+        for key in self._dict:
+            yield (key, self[key])
+
+    def pop(self, key: K) -> tuple["FrozenDict[K, V]", V]:
+        value = self[key]
+        new_dict = dict(self._dict)
+        new_dict.pop(key)
+        new_self = type(self)(new_dict)
+        return new_self, value
+
+    def unfreeze(self) -> dict[K, V]:
+        return unfreeze(self)
+
+    def tree_flatten_with_keys(self) -> tuple[tuple[tuple[jax.tree_util.DictKey, Any], ...], tuple[K, ...]]:
+        sorted_keys = sorted(self._dict)
+        return tuple([(jax.tree_util.DictKey(k), self._dict[k]) for k in sorted_keys]), tuple(sorted_keys)
+
+    @classmethod
+    def tree_unflatten(cls, keys: tuple[K, ...], values: tuple[Any, ...]) -> "FrozenDict[K, V]":
+        return cls({k: v for k, v in zip(keys, values)}, __unsafe_skip_copy__=True)
+
+
+def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: ANN401
+    if isinstance(x, FrozenDict):
+        return jax.tree_util.tree_map(lambda y: y, x._dict)
+    elif isinstance(x, dict):
+        ys = {}
+        for key, value in x.items():
+            ys[key] = unfreeze(value)
+        return ys
+    else:
+        return x

xax/utils/types/hashable_array.py

@@ -0,0 +1,31 @@
+"""Defines a hashable array wrapper.
+
+Since Jax relies extensively on hashing, and we sometimes want to treat Jax
+arrays as constants, this wrapper lets us ensure that Jax and Numpy arrays can
+be hashed for Jitting.
+"""
+
+import jax.numpy as jnp
+import numpy as np
+
+
+class HashableArray:
+    def __init__(self, array: np.ndarray | jnp.ndarray) -> None:
+        if not isinstance(array, (np.ndarray, jnp.ndarray)):
+            raise ValueError(f"Expected np.ndarray or jnp.ndarray, got {type(array)}")
+        self.array = array
+        self._hash: int | None = None
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            self._hash = hash(self.array.tobytes())
+        return self._hash
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, HashableArray):
+            return False
+        return bool(jnp.array_equal(self.array, other.array))
+
+
+def hashable_array(array: np.ndarray | jnp.ndarray) -> HashableArray:
+    return HashableArray(array)

{xax-0.1.6.dist-info → xax-0.1.7.dist-info}/METADATA

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

{xax-0.1.6.dist-info → xax-0.1.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=rjqydWhxQVUAj3lXgFpzj4iLFOdDJGHArfAH7_QSkhk,13504
+xax/__init__.py,sha256=Yxbt78w3k1Y827EjR_kiT6_01FyRrJKn1U5WkC5I2Ik,13453
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -7,8 +7,8 @@ xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
 xax/core/state.py,sha256=y123fL7pMgk25TPG6KN0LRIF_eYnD9eP7OfqtoQJGNE,2178
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
-xax/nn/equinox.py,sha256=1Ck6ycz76dhit2LHX4y2lp3WJSPsDuRt7TK7AxxQhww,4837
-xax/nn/export.py,sha256=bu2m-4FDnadEhXDb9zM6SgOZvsf5p4xiee1sFZyNF7c,5510
+xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
+xax/nn/export.py,sha256=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
 xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
 xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
@@ -31,7 +31,7 @@ xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCn
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
 xax/task/mixins/checkpointing.py,sha256=a6tVyISsDIz68rrhb1rAh3rjQlqkDVJCmSBmETQrnRM,8480
-xax/task/mixins/compile.py,sha256=9pVJEUvizu6-6tq0HaMtHGNSi9Yk_mxNyqBFimcfwL0,3683
+xax/task/mixins/compile.py,sha256=8jEdlGs-a14N_CwZA3Rxe461MT83dyIDr3Z56VkjviQ,3693
 xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
 xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
 xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
@@ -43,7 +43,7 @@ xax/task/mixins/train.py,sha256=vsH_QpyrThlh9AzWnyvDJv58Y8U_516oi8gmMq_0iMg,2233
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
-xax/utils/jax.py,sha256=eObvWt2DraCs2IMDZSdQ0rRk8tA3P5XBlF_UeVq7Aro,5480
+xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
@@ -53,8 +53,11 @@ xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,1705
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
-xax-0.1.6.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.6.dist-info/METADATA,sha256=vKxhuOt02ALjFV9fAt-rPVTwvqX4uNr_shL1DGEotA4,1877
-xax-0.1.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.6.dist-info/RECORD,,
+xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
+xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
+xax-0.1.7.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.7.dist-info/METADATA,sha256=KWlXilX8eGmFEi9-FR-QyX3pyUEEbT2dCjx8tFA1yKU,1877
+xax-0.1.7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.7.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.7.dist-info/RECORD,,