Merge branch 'main' into intersect-axes

README.md

@@ -47,6 +47,9 @@ Head = hax.Axis("head", 8)  # number of attention heads
 Key = hax.Axis("key", 64)  # key size
 Embed = hax.Axis("embed", 512)  # embedding size
 
+# alternatively:
+#Pos, KPos, Head, Key, Embed = hax.make_axes(pos=1024, key_pos=1024, head=8, key=64, embed=512)
+
 
 def attention_scores(Key, KPos, query, key, mask):
     # how similar is each query to each key

docs/api.md

@@ -27,6 +27,7 @@ Occasionally, an axis size can be inferred in some circumstances but not others.
 
 ### Axis Manipulation
 
+::: haliax.make_axes
 ::: haliax.axis.axis_name
 ::: haliax.axis.concat_axes
 ::: haliax.axis.union_axes

src/haliax/__init__.py

@@ -32,6 +32,7 @@
     ds,
     dslice,
     eliminate_axes,
+    make_axes,
     selects_axis,
 )
 from .core import (
@@ -893,6 +894,9 @@ def true_divide(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     "AxisSpec",
     "AxisSelection",
     "AxisSelector",
+    "make_axes",
+    "axis_name",
+    "axis_size",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",

src/haliax/axis.py

@@ -28,6 +28,17 @@ def __str__(self):
         return f"{self.name}({self.size})"
 
 
+def make_axes(**kwargs: int) -> Tuple[Axis, ...]:
+    """
+    Convenience function for creating a tuple of Axis objects.
+
+    Example:
+    ```
+    X, Y = axes(X=10, Y=20)
+    """
+    return tuple(Axis(name, size) for name, size in kwargs.items())
+
+
 AxisSelector = Union[Axis, str]
 """AxisSelector is a type that can be used to select a single axis from an array. str or Axis"""
 AxisSelection = Union[AxisSelector, Sequence[AxisSelector]]

src/haliax/core.py

@@ -373,11 +373,10 @@ def __getitem__(self, idx: SliceSpec) -> "NamedArray":
 
         Supports indexing like:
 
-        >>> X = Axis("x", 10)
-        >>> Y = Axis("y", 20)
+        >>> X, Y = haliax.make_axes(X=10, Y=20)
         >>> arr = haliax.random.randint(jax.random.PRNGKey(0), (X, Y), 0, X.size)
         # slice with ints or slices
-        >>> arr[{"x": 1, "y": slice(0,10,new_axis=2)}]
+        >>> arr[{"x": 1, "y": slice(0,10,2)}]
         >>> Z = Axis("z", 3)
         # so-called "advanced indexing" with NamedArrays.
         >>> index_arr = NamedArray(np.array([1, 2, 3]), Z)

src/haliax/nn/mlp.py

@@ -44,6 +44,7 @@ def init(
         depth: int,
         activation: Callable = relu,
         *,
+        out_first: bool = True,
         use_bias: bool = True,
         use_final_bias: bool = True,
         key: PRNGKeyArray,
@@ -57,36 +58,34 @@ def init(
 
         layers = []
 
+        kwargs: dict = {
+            "use_bias": use_bias,
+            "dot_general": dot_general,
+            "init_scale": init_scale,
+            "out_first": out_first,
+        }
+
+        last_kwargs: dict = {
+            "use_bias": use_final_bias,
+            "dot_general": dot_general,
+            "init_scale": init_scale,
+            "out_first": out_first,
+        }
+
         if depth == 0:
             # special case: no hidden layers
-            layers.append(
-                Linear.init(
-                    Input, Output, use_bias=use_final_bias, key=keys[0], dot_general=dot_general, init_scale=init_scale
-                )
-            )
+            layers.append(Linear.init(Input, Output, key=keys[0], **last_kwargs))
         else:
             # first hidden layer
-            layers.append(
-                Linear.init(
-                    Input, Width, use_bias=use_bias, key=keys[0], dot_general=dot_general, init_scale=init_scale
-                )
-            )
+            layers.append(Linear.init(Input, Width, key=keys[0], **kwargs))
             # middle hidden layers
             cur = Width
             next = Width2
             for i in range(1, depth):
-                layers.append(
-                    Linear.init(
-                        cur, next, use_bias=use_bias, key=keys[i], dot_general=dot_general, init_scale=init_scale
-                    )
-                )
+                layers.append(Linear.init(cur, next, key=keys[i], **kwargs))
                 cur, next = next, cur
-            # final hidden layer
-            layers.append(
-                Linear.init(
-                    cur, Output, use_bias=use_final_bias, key=keys[-1], dot_general=dot_general, init_scale=init_scale
-                )
-            )
+            # final layer
+            layers.append(Linear.init(cur, Output, key=keys[-1], **last_kwargs))
 
         return MLP(
             layers=tuple(layers),

src/haliax/nn/scan.py

@@ -300,7 +300,7 @@ def unstacked(self) -> Sequence[M]:
         """
 
         def unbatch_leaf(x):
-            if haliax.is_named_array(x):
+            if isinstance(x, haliax.core.NamedArray):
                 if haliax.selects_axis(x.axes, self.Block):
                     return haliax.unbind(x, self.Block)
                 else: