haliax 1.4.dev325__tar.gz → 1.4.dev326__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev325
+Version: 1.4.dev326
 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/

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

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

{haliax-1.4.dev325 → haliax-1.4.dev326}/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.dev326}/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.dev326}/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.dev326}/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/src/haliax/__about__.py

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