haliax 1.4.dev325__tar.gz → 1.4.dev327__tar.gz

{haliax-1.4.dev325 → haliax-1.4.dev327}/PKG-INFO

@@ -1,11 +1,12 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev325
+Version: 1.4.dev327
 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.dev327/src/haliax/__about__.py

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

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/__init__.py

@@ -34,6 +34,8 @@ from .axis import (
     dslice,
     eliminate_axes,
     make_axes,
+    replace_axis,
+    resolve_axis,
     selects_axis,
 )
 from .core import (
@@ -1060,6 +1062,8 @@ __all__ = [
     "stack",
     "concatenate",
     "eliminate_axes",
+    "resolve_axis",
+    "replace_axis",
     "selects_axis",
     "concat_axes",
     "concat_axis_specs",

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/_src/state_dict.py

@@ -15,7 +15,8 @@ from jaxtyping import PyTree
 
 import haliax.partitioning as partitioning
 from haliax._src.util import index_where
-from haliax.core import NamedArray, named
+from haliax.axis import Axis
+from haliax.core import NamedArray, flatten_axes, named
 from haliax.jax_utils import is_jax_array_like, is_scalarish
 
 
@@ -390,24 +391,22 @@ def flatten_linear_layers(tree: T) -> T:
         weight = layer.weight
         bias = layer.bias
 
+        new_Out: Axis = flatten_axes(layer.Out, "__OUT__")
+        new_In: Axis = flatten_axes(layer.In, "__IN__")
+
         if weight.array is not None:
             out_first = layer.out_first
-            weight = weight.flatten_axes(layer.Out, "__OUT__").flatten_axes(layer.In, "__IN__")
+            weight = weight.flatten_axes(layer.Out, new_Out).flatten_axes(layer.In, new_In)
 
             if out_first:
                 weight = weight.rearrange((..., "__OUT__", "__IN__"))
             else:
                 weight = weight.rearrange((..., "__IN__", "__OUT__"))
 
-            if bias is not None:
-                bias = bias.flatten_axes(layer.Out, "__OUT__")
-
-            In = weight.resolve_axis("__IN__")
-            Out = weight.resolve_axis("__OUT__")
+        if isinstance(bias, NamedArray):
+            bias = bias.flatten_axes(layer.Out, new_Out)
 
-            return dataclasses.replace(layer, weight=weight, bias=bias, In=In, Out=Out)  # type: ignore
-        else:
-            return layer
+        return dataclasses.replace(layer, weight=weight, bias=bias, In=new_In, Out=new_Out)  # type: ignore
 
     return jax.tree.map(_flatten_linear, tree, is_leaf=lambda x: isinstance(x, Linear))
 
@@ -438,7 +437,7 @@ def unflatten_linear_layers(template: T, tree_with_flattened_linears: T) -> T:
             weight = weight.unflatten_axis("__OUT__", template.Out).unflatten_axis("__IN__", template.In)
             weight = weight.rearrange(template.weight.axes)
 
-        if bias is not None:
+        if isinstance(bias, NamedArray):
             bias = bias.unflatten_axis("__OUT__", template.Out)
             assert template.bias is not None, "Flattened bias but template has no bias"
             bias = bias.rearrange(template.bias.axes)

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/axis.py

@@ -394,6 +394,52 @@ def axis_size(ax: AxisSpec) -> int:
         return prod(axis.size for axis in ensure_tuple(ax))  # type: ignore
 
 
+@typing.overload
+def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelector) -> Axis:
+    ...
+
+
+@typing.overload
+def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec:
+    ...
+
+
+def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec:
+    """
+    Returns the axis or axes in axis_spec that match the name of axis_selection.
+
+    If axis_selection is a str or axis, returns a single Axis. If it is a sequence, returns a sequence of Axes.
+
+    If an axis is present with a different size, raises ValueError.
+    """
+    ax: Axis
+    if isinstance(axis_selection, str | Axis):
+        name = axis_name(axis_selection)
+        for ax in ensure_tuple(axis_spec):
+            if axis_name(ax) == name:
+                if isinstance(axis_selection, Axis) and axis_size(ax) != axis_size(axis_selection):
+                    raise ValueError(f"Axis {name} has different sizes in {axis_spec} and {axis_selection}")
+                return ax
+        raise ValueError(f"Axis {name} not found in {axis_spec}")
+    else:
+        as_map = axis_spec_to_shape_dict(axis_spec)
+        out: list[Axis] = []
+
+        for ax in ensure_tuple(axis_selection):  # type: ignore
+            name = axis_name(ax)
+            if name not in as_map:
+                raise ValueError(f"Axis {name} not found in {axis_spec}")
+            if isinstance(ax, Axis):
+                if as_map[name] != ax.size:  # type: ignore
+                    raise ValueError(f"Axis {name} has different sizes in {axis_spec} and {axis_selection}")
+                else:
+                    out.append(ax)
+            else:
+                out.append(Axis(name, as_map[name]))
+
+        return tuple(out)
+
+
 class dslice(eqx.Module):
     """
     Dynamic slice, comprising a (start, length) pair. Also aliased as ds.
@@ -576,6 +622,7 @@ __all__ = [
     "is_axis_compatible",
     "overlapping_axes",
     "replace_axis",
+    "resolve_axis",
     "selects_axis",
     "union_axes",
     "without_axes",

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/core.py

@@ -1143,12 +1143,65 @@ def rename(array: NamedArray, renames: Mapping[AxisSelector, AxisSelector]) -> N
     return NamedArray(array.array, new_axes)
 
 
+@typing.overload
+def flatten_axes(axis: Axis, old_axes: Axis, new_axis: AxisSelector) -> Axis:
+    pass
+
+
+@typing.overload
+def flatten_axes(axis: AxisSpec, new_axis: AxisSelector) -> Axis:
+    pass
+
+
+@typing.overload
+def flatten_axes(axis: AxisSpec, old_axes: AxisSelection, new_axis: AxisSelector) -> AxisSpec:
+    pass
+
+
+@typing.overload
 def flatten_axes(array: NamedArray, old_axes: AxisSelection, new_axis: AxisSelector) -> NamedArray:
+    pass
+
+
+def flatten_axes(  # type: ignore
+    # array: NamedArray | AxisSpec, old_axes: AxisSelection, new_axis: AxisSelector
+    *args,
+    **kwargs,
+) -> NamedArray | AxisSpec:
     """
     Merge a sequence of axes into a single axis. The new axis must have the same size as the product of the old axes.
 
-    The new axis is always inserted starting at the index of the first old axis in theunderlying array.
+    The new axis is always inserted starting at the index of the first old axis in the underlying array.
+
+    This function can be used in two ways:
+
+    * `flatten_axes(array, old_axes, new_axis)`: merge the old axes of the array into a new axis
+    * `flatten_axes(axes, old_axes, new_axis)`: merge the old axes into a new axis
     """
+
+    if len(args) + len(kwargs) == 2:
+        return _simple_flatten(*args, **kwargs)
+    else:
+        return _full_flatten(*args, **kwargs)
+
+
+def _simple_flatten(axis: AxisSpec, new_axis: AxisSelector) -> Axis:
+    size = haliax.axis_size(axis)
+    if isinstance(new_axis, Axis):
+        if new_axis.size != size:
+            raise ValueError(f"Cannot merge {axis} into {new_axis}: size mismatch")
+        return new_axis
+
+    assert isinstance(new_axis, str)
+    return Axis(new_axis, size)
+
+
+def _full_flatten(
+    array: NamedArray | AxisSpec, old_axes: AxisSelection, new_axis: AxisSelector
+) -> NamedArray | AxisSpec:
+    if isinstance(array, str | Axis | Sequence):
+        return _flatten_axis_spec(array, old_axes, new_axis)
+
     old_axes = ensure_tuple(old_axes)
     old_axes = array.resolve_axis(old_axes)
     total_axis_size = haliax.axis_size(old_axes)
@@ -1187,6 +1240,42 @@ def flatten_axes(array: NamedArray, old_axes: AxisSelection, new_axis: AxisSelec
     return NamedArray(raw_array, tuple(new_axes))
 
 
+def _flatten_axis_spec(axes: AxisSpec, old_axes: AxisSelection, new_axis: AxisSelector) -> AxisSpec:
+    axes = ensure_tuple(axes)
+    old_axes = ensure_tuple(old_axes)
+    old_axes = haliax.axis.resolve_axis(axes, old_axes)
+    total_axis_size = haliax.axis_size(old_axes)
+
+    if isinstance(new_axis, Axis):
+        if new_axis.size != total_axis_size:
+            raise ValueError(f"Cannot merge {old_axes} into {new_axis}: size mismatch")
+    else:
+        assert isinstance(new_axis, str)
+        new_axis = Axis(new_axis, total_axis_size)
+
+    if len(old_axes) == 0:  # type: ignore
+        return (new_axis,) + axes
+
+    # ensure that the old_axes are contiguous
+    # we basically ensure that the old_axes occur after the index of the first old_axis
+    intermediate_axes: List[Axis] = []
+    new_axes: List[Axis] = []
+    index_of_first_old_axis = None
+    for i, ax in enumerate(axes):
+        if ax in old_axes:  # type: ignore
+            if index_of_first_old_axis is None:
+                index_of_first_old_axis = i
+                intermediate_axes.extend(old_axes)  # type: ignore
+                new_axes.append(new_axis)
+            else:
+                continue
+        else:
+            intermediate_axes.append(ax)
+            new_axes.append(ax)
+
+    return tuple(new_axes)
+
+
 def unflatten_axis(array: NamedArray, axis: AxisSelector, new_axes: AxisSpec) -> NamedArray:
     """
     Split an axis into a sequence of axes. The old axis must have the same size as the product of the new axes.

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/partitioning.py

@@ -8,7 +8,7 @@ from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Seque
 
 import equinox as eqx
 import jax
-from equinox import module_update_wrapper
+from equinox import is_array, module_update_wrapper
 from jax.lax import with_sharding_constraint
 from jax.sharding import Mesh, NamedSharding, PartitionSpec, SingleDeviceSharding
 from jaxtyping import PyTree
@@ -20,7 +20,7 @@ from .axis import Axis, AxisSelection, AxisSelector
 from .core import NamedArray
 from .jax_utils import Static, is_in_jit, is_jax_array_like, is_on_mac_metal
 from .tree_util import hashable_combine, hashable_partition
-from .util import StringHolderEnum, ensure_tuple, is_named_array
+from .util import StringHolderEnum, ensure_tuple
 
 
 PhysicalAxisSpec = Union[(str), Sequence[str]]
@@ -274,7 +274,7 @@ class _NamedJitWrapper(eqx.Module):
         if out_axis_resources is None:
             out_axis_resources = axis_resources
 
-        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_jax_array_like)
+        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_array)
         dynamic = (self._dynamic_fun, dynamic_argspec)
 
         donate_args = self._donate_args
@@ -436,7 +436,7 @@ def named_jit(
             **pjit_args,
         )
 
-    dynamic_fun, static_fun = hashable_partition(fn, is_jax_array_like)
+    dynamic_fun, static_fun = hashable_partition(fn, is_array)
 
     wrapper = _NamedJitWrapper(
         fn,
@@ -514,7 +514,7 @@ def _named_pjit_cache(fun_names, **jitkwargs) -> WrappedCallable:
         fun = hashable_combine(dynamic_fun, static_fun)
         args, kwargs = hashable_combine(dynamic_spec, static_spec)
         out = fun(*args, **kwargs)
-        out_dynamic, out_static = hashable_partition(out, is_jax_array_like)
+        out_dynamic, out_static = hashable_partition(out, is_array)
         return out_dynamic, Static(out_static)
 
     fun_name, fun_qualname = fun_names
@@ -543,7 +543,7 @@ def _cached_filter_eval_shape(fun, *args, **kwargs):
     eval_shape is surprisingly expensive, so we cache it. We use this for named_pjit for evaluating resource partitions
     of the output.
     """
-    dynamic, static = hashable_partition((fun, args, kwargs), is_jax_array_like)
+    dynamic, static = hashable_partition((fun, args, kwargs), is_array)
     if static not in _eval_shape_cache:
         _eval_shape_cache[static] = eqx.filter_eval_shape(fun, *args, **kwargs)
 

{haliax-1.4.dev325 → haliax-1.4.dev327}/src/haliax/quantization.py

@@ -10,7 +10,6 @@ from typing import Optional, Protocol, TypeVar
 import equinox as eqx
 import jax
 from jax import numpy as jnp
-from jax._src.tree_util import BuiltInKeyEntry
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
 from jax.typing import DTypeLike
 
@@ -253,7 +252,7 @@ def _matches_target_fp8(key_path, config: Fp8Config) -> bool:
     return re.match(config.targets, key_path_str) is not None
 
 
-def _key_path_to_str(key_path: tuple[BuiltInKeyEntry, ...]) -> str:
+def _key_path_to_str(key_path: tuple) -> str:
     out = ""
     for k in key_path:
         match k:

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

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