xax 0.1.6__tar.gz → 0.1.8__tar.gz

{xax-0.1.6/xax.egg-info → xax-0.1.8}/PKG-INFO

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

{xax-0.1.6 → xax-0.1.8}/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.8"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -97,8 +97,6 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
-    "HashableArray",
-    "hashable_array",
     "jit",
     "save_jaxpr_dot",
     "ColoredFormatter",
@@ -132,6 +130,9 @@ __all__ = [
     "snakecase_to_camelcase",
     "uncolored",
     "wrapped",
+    "FrozenDict",
+    "HashableArray",
+    "hashable_array",
 ]
 
 __all__ += [
@@ -161,7 +162,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).
@@ -253,8 +254,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 +287,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.
@@ -408,7 +410,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 +450,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-0.1.6 → xax-0.1.8}/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-0.1.6 → xax-0.1.8}/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-0.1.6 → xax-0.1.8}/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-0.1.6 → xax-0.1.8}/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-0.1.8/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-0.1.8/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 → xax-0.1.8/xax.egg-info}/PKG-INFO

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

{xax-0.1.6 → xax-0.1.8}/xax.egg-info/SOURCES.txt

@@ -63,4 +63,7 @@ xax/utils/pytree.py
 xax/utils/tensorboard.py
 xax/utils/text.py
 xax/utils/data/__init__.py
-xax/utils/data/collate.py
+xax/utils/data/collate.py
+xax/utils/types/__init__.py
+xax/utils/types/frozen_dict.py
+xax/utils/types/hashable_array.py