haliax 1.4.dev322__tar.gz → 1.4.dev324__tar.gz

{haliax-1.4.dev322 → haliax-1.4.dev324}/PKG-INFO

@@ -1,12 +1,11 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev322
+Version: 1.4.dev324
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/
 Project-URL: Documentation, https://haliax.readthedocs.io/en/latest/
 Author-email: David Hall <dlwh@cs.stanford.edu>
-License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License

haliax-1.4.dev324/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev324"

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/_src/state_dict.py

@@ -16,7 +16,7 @@ from jaxtyping import PyTree
 import haliax.partitioning as partitioning
 from haliax._src.util import index_where
 from haliax.core import NamedArray, named
-from haliax.jax_utils import is_jax_array_like
+from haliax.jax_utils import is_jax_array_like, is_scalarish
 
 
 try:
@@ -166,7 +166,9 @@ def to_state_dict(tree: PyTree, prefix: Optional[str] = None) -> StateDict:
     Returns:
         The state dict representation of the input tree.
     """
-    if isinstance(tree, eqx.Module):
+    if tree is None:
+        return {}
+    elif isinstance(tree, eqx.Module):
         if hasattr(tree, "to_state_dict"):
             state_dict = tree.to_state_dict(prefix)
         else:
@@ -198,6 +200,11 @@ def to_state_dict(tree: PyTree, prefix: Optional[str] = None) -> StateDict:
                 state_dict = {}
         else:
             raise ValueError("Cannot convert a leaf value to a state dict without a prefix")
+    elif is_scalarish(tree):
+        if prefix is not None:
+            state_dict = {prefix: tree}
+        else:
+            raise ValueError("Cannot convert a leaf value to a state dict without a prefix")
     else:
         raise ValueError(f"Unsupported type {type(tree)}")
 
@@ -291,37 +298,46 @@ def to_numpy_state_dict(model, prefix: Optional[str] = None) -> StateDict:
     with jax.default_device(jax.local_devices(backend="cpu")[0]):
 
         def get_to_cpu(arr):
-            if not is_jax_array_like(arr):
+            if is_scalarish(arr):
+                return arr
+            elif not is_jax_array_like(arr):
                 return arr
             elif isinstance(arr, np.ndarray):
                 return arr
-            elif arr.is_fully_addressable:
-                r = np.array(arr)
-                return r
-            else:
-                # unfortunately, jax's allgather seems to replicate to every device rather than every host
-                # which doesn't work for ~7B parameter models on TPU (assuming we also have optimizer state)
-                # this approach limits us to <64B parameters, but that's good enough for now
-                # we're going to do something a bit fancy, where we shard the model into a (process, device) mesh,
-                # then look for some axis along which we can shard the array, and then we'll do an allgather
-                # via pjit. If we can't find one, we'll just fully replicate since it probably isn't that big.
-                # TODO: ensure that this mesh arranges devices correctly
-                # (jax seems to do this internally itself, so we should be fine?)
-                process_mesh = Mesh(np.array(jax.devices()).reshape((jax.process_count(), -1)), ("process", "device"))
-                # now we need to find an axis along which we can shard the array.
-                # for this, we need to find an axis s.t. size(axis) % local_devices == 0
-
-                try:
-                    axis_to_shard = index_where(
-                        lambda axis_size: axis_size % process_mesh.devices.size == 0, arr.shape
+            elif is_jax_array_like(arr):
+                if arr.is_fully_addressable:
+                    r = np.array(arr)
+                    return r
+                else:
+                    # unfortunately, jax's allgather seems to replicate to every device rather than every host
+                    # which doesn't work for ~7B parameter models on TPU (assuming we also have optimizer state)
+                    # this approach limits us to <64B parameters, but that's good enough for now
+                    # we're going to do something a bit fancy, where we shard the model into a (process, device) mesh,
+                    # then look for some axis along which we can shard the array, and then we'll do an allgather
+                    # via pjit. If we can't find one, we'll just fully replicate since it probably isn't that big.
+                    # TODO: ensure that this mesh arranges devices correctly
+                    # (jax seems to do this internally itself, so we should be fine?)
+                    process_mesh = Mesh(
+                        np.array(jax.devices()).reshape((jax.process_count(), -1)), ("process", "device")
                     )
-                except ValueError:
-                    return np.array(arr)
-
-                shardings = [None if i != axis_to_shard else "device" for i in range(len(arr.shape))]
-                sharding = NamedSharding(process_mesh, PartitionSpec(*shardings))
-                out = jax.jit(lambda x: x, out_shardings=sharding)(arr)
-                return np.array(out)
+                    # now we need to find an axis along which we can shard the array.
+                    # for this, we need to find an axis s.t. size(axis) % local_devices == 0
+
+                    try:
+                        axis_to_shard = index_where(
+                            lambda axis_size: axis_size % process_mesh.devices.size == 0, arr.shape
+                        )
+                    except ValueError:
+                        return np.array(arr)
+
+                    shardings = [None if i != axis_to_shard else "device" for i in range(len(arr.shape))]
+                    sharding = NamedSharding(process_mesh, PartitionSpec(*shardings))
+                    out = jax.jit(lambda x: x, out_shardings=sharding)(arr)
+                    return np.array(out)
+            elif is_scalarish(arr):
+                return np.asarray(arr)
+            else:
+                raise ValueError(f"Unsupported type {type(arr)}")
 
         # need to make sure the model is on *this machine* and *this machine's CPU* before saving
         model = jax.tree.map(lambda arr: get_to_cpu(arr), model)

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/core.py

@@ -246,7 +246,9 @@ class NamedArray:
                 try:
                     axis_index = index_where(lambda a: a.name == ax_name, self.axes)
                     if axis_index >= 0:
-                        warnings.warn("Found axis with same name but different size.", UserWarning)
+                        raise RuntimeError(
+                            f"Found axis with same name but different size: {axis} vs {self.axes[axis_index]}"
+                        )
                     return axis_index
                 except ValueError:
                     return None

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/jax_utils.py

@@ -152,7 +152,7 @@ def is_scalarish(x):
     if isinstance(x, haliax.NamedArray):
         return x.ndim == 0
     else:
-        return jnp.isscalar(x) or x.shape == ()
+        return jnp.isscalar(x) or (getattr(x, "shape", None) == ())
 
 
 def is_on_mac_metal():

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/nn/scan.py

@@ -8,7 +8,8 @@ from jax import numpy as jnp
 
 import haliax
 import haliax.util
-from haliax.jax_utils import filter_checkpoint
+from haliax.jax_utils import filter_checkpoint, is_jax_array_like
+from haliax.util import is_jax_or_hax_array_like
 
 from .._src.state_dict import ModuleWithStateDictSerialization, StateDict, with_prefix
 from ..axis import Axis
@@ -167,7 +168,9 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
         state_dict: StateDict = {}
         for i, block in enumerate(self.blocks):
             my_prefix = with_prefix(prefix, str(i))
-            state_dict.update(block.to_state_dict(my_prefix))
+            # we can't assume to_state_dict is implemented, so we have to do it manually
+            block_dict = haliax.state_dict.to_state_dict(block, my_prefix)
+            state_dict.update(block_dict)
 
         return state_dict
 
@@ -413,7 +416,7 @@ def _unstack_state_dict(state_dict: StateDict, prefix: Optional[str] = None) ->
     assert prefix is not None
 
     for k, v in state_dict.items():
-        if k.startswith(prefix) and v is not None:
+        if k.startswith(prefix) and is_jax_or_hax_array_like(v):
             for i, v_i in enumerate(v):
                 new_dict[f"{prefix}{i}.{k[len(prefix):]}"] = v_i
         else:

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/state_dict.py

@@ -1,5 +1,10 @@
 from typing import Optional, TypeVar
 
+import equinox
+
+from haliax.jax_utils import is_jax_array_like
+from haliax.types import FilterSpec
+
 from ._src.state_dict import (
     ModuleWithStateDictSerialization,
     StateDict,
@@ -18,12 +23,22 @@ from ._src.state_dict import (
 T = TypeVar("T")
 
 
-def to_torch_compatible_state_dict(t: T, *, flatten_linear: bool = True, prefix: Optional[str] = None) -> StateDict:
+def to_torch_compatible_state_dict(
+    t: T, *, flatten_linear: bool = True, prefix: Optional[str] = None, filter: FilterSpec = is_jax_array_like
+) -> StateDict:
     """
     Convert a tree to a state dict that is compatible with torch-style state dicts.
 
     This applies [haliax.state_dict.flatten_linear_layers][] followed by [haliax.state_dict.to_state_dict][]
+
+    Args:
+        t: The tree to convert
+        flatten_linear: Whether to flatten linear layers
+        prefix: The prefix to use for the state dict keys
+        filter: The filter to use for selecting which nodes to include in the state dict. By default, this includes only
+            array-like objects (e.g. JAX and NumPy arrays).
     """
+    t = equinox.filter(t, filter)
     if flatten_linear:
         t = flatten_linear_layers(t)
     return to_numpy_state_dict(t, prefix=prefix)

{haliax-1.4.dev322 → haliax-1.4.dev324}/src/haliax/types.py

@@ -1,8 +1,9 @@
-from typing import Any, Literal, Protocol, Tuple, TypeAlias, Union
+from typing import Any, Callable, Literal, Protocol, Tuple, TypeAlias, Union
 
 import jax.numpy as jnp
 import numpy as np
 from jax.lax import Precision
+from jaxtyping import PyTree
 
 
 DType: TypeAlias = np.dtype
@@ -30,3 +31,13 @@ IntScalar = Union[int, jnp.ndarray]
 PrecisionLike = Union[None, str, Precision, Tuple[str, str], Tuple[Precision, Precision]]
 
 GatherScatterModeStr = Literal["promise_in_bounds", "clip", "drop", "fill"]
+
+
+FilterSpec = Union[bool, Callable[[Any], bool]]
+"""
+A filter specification. Typically used on a pytree to filter out certain subtrees. Boolean values are
+treated as-is, while callables are called on each element of the pytree. If the callable returns True, the element
+is kept, otherwise it is filtered out.
+"""
+
+FilterTree = FilterSpec | PyTree[FilterSpec]

haliax-1.4.dev322/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev322"