haliax 1.4.dev315__tar.gz → 1.4.dev317__tar.gz

{haliax-1.4.dev315 → haliax-1.4.dev317}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev315
+Version: 1.4.dev317
 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.dev315 → haliax-1.4.dev317}/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

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

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

{haliax-1.4.dev315 → haliax-1.4.dev317}/src/haliax/__init__.py

@@ -32,6 +32,7 @@ from .axis import (
     ds,
     dslice,
     eliminate_axes,
+    make_axes,
     selects_axis,
 )
 from .core import (
@@ -893,6 +894,9 @@ __all__ = [
     "AxisSpec",
     "AxisSelection",
     "AxisSelector",
+    "make_axes",
+    "axis_name",
+    "axis_size",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",

{haliax-1.4.dev315 → haliax-1.4.dev317}/src/haliax/axis.py

@@ -28,6 +28,17 @@ class Axis:
         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]]

{haliax-1.4.dev315 → haliax-1.4.dev317}/src/haliax/core.py

@@ -373,11 +373,10 @@ class 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)

{haliax-1.4.dev315 → haliax-1.4.dev317}/src/haliax/nn/mlp.py

@@ -44,6 +44,7 @@ class MLP(eqx.Module):
         depth: int,
         activation: Callable = relu,
         *,
+        out_first: bool = True,
         use_bias: bool = True,
         use_final_bias: bool = True,
         key: PRNGKeyArray,
@@ -57,36 +58,34 @@ class MLP(eqx.Module):
 
         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),

{haliax-1.4.dev315 → haliax-1.4.dev317}/src/haliax/nn/scan.py

@@ -300,7 +300,7 @@ class Stacked(eqx.Module, Generic[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:

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/core_test.py

@@ -9,9 +9,7 @@ from haliax import Axis, NamedArray
 
 
 def test_unary_np_functions():
-    Height = Axis("Height", 2)
-    Width = Axis("Width", 3)
-    Depth = Axis("Depth", 4)
+    Height, Width, Depth = hax.make_axes(Height=2, Width=3, Depth=4)
 
     m1 = NamedArray(jnp.ones((Height.size, Width.size, Depth.size)), (Height, Width, Depth))
 
@@ -29,9 +27,7 @@ def test_unary_np_functions():
 
 
 def test_reduction_functions():
-    Height = Axis("Height", 2)
-    Width = Axis("Width", 3)
-    Depth = Axis("Depth", 4)
+    Height, Width, Depth = hax.make_axes(Height=2, Width=3, Depth=4)
 
     rand_m = jax.random.uniform(PRNGKey(0), (Height.size, Width.size, Depth.size))
 
@@ -69,9 +65,7 @@ def test_reduction_functions():
 
 
 def test_reduction_functions_with_where():
-    H = Axis("H", 2)
-    W = Axis("W", 3)
-    D = Axis("D", 4)
+    H, W, D = hax.make_axes(H=2, W=3, D=4)
 
     rand_m = jax.random.uniform(PRNGKey(0), (H.size, W.size, D.size))
 
@@ -118,9 +112,7 @@ def test_reduction_functions_with_where():
 
 
 def test_split():
-    Height = Axis("Height", 2)
-    Width = Axis("Width", 3)
-    Depth = Axis("Depth", 4)
+    Height, Width, Depth = hax.make_axes(Height=2, Width=3, Depth=4)
 
     D10 = Axis("Depth", Depth.size * 10)
 
@@ -144,9 +136,7 @@ def test_split():
 
 
 def test_take():
-    Height = Axis("Height", 2)
-    Width = Axis("Width", 3)
-    Depth = Axis("Depth", 4)
+    Height, Width, Depth = hax.make_axes(Height=2, Width=3, Depth=4)
     named1 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
 
     assert jnp.all(jnp.equal(hax.take(named1, Height, 0).array, named1.array[0]))
@@ -178,9 +168,7 @@ def test_take():
 
 
 def test_take_overlapping_names():
-    Height = Axis("Height", 20)
-    Width = Axis("Width", 30)
-    Depth = Axis("Depth", 40)
+    Height, Width, Depth = hax.make_axes(Height=20, Width=30, Depth=40)
     named1 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
 
     Height2 = Axis("Height", 10)
@@ -198,9 +186,7 @@ def test_take_overlapping_2():
     def cross_entropy(logits: hax.NamedArray, labels: hax.NamedArray) -> hax.NamedArray:
         return hax.take(logits, Embed, labels)  # extract log probability of the correct token
 
-    Embed = Axis("Embed", 10)
-    Block = Axis("Block", 20)
-    Batch = Axis("Batch", 30)
+    Embed, Block, Batch = hax.make_axes(Embed=10, Block=20, Batch=30)
     logits = hax.random.uniform(PRNGKey(0), (Batch, Block, Embed))
     labels = hax.random.randint(PRNGKey(0), (Batch, Block), 0, Embed.size)
 
@@ -223,9 +209,7 @@ def test_take_overlapping_2():
 
 
 def test_cumsum_etc():
-    Height = Axis("Height", 2)
-    Width = Axis("Width", 3)
-    Depth = Axis("Depth", 4)
+    Height, Width, Depth = hax.make_axes(Height=2, Width=3, Depth=4)
 
     named1 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
 
@@ -255,10 +239,7 @@ def test_cumsum_etc():
 
 
 def test_rearrange():
-    H = Axis("H", 2)
-    W = Axis("W", 3)
-    D = Axis("D", 4)
-    C = Axis("C", 5)
+    H, W, D, C = hax.make_axes(H=2, W=3, D=4, C=5)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D, C))
 
@@ -304,9 +285,7 @@ def test_rearrange():
 
 def test_rearrange_unused_ellipsis():
     # Make sure we just ignore the ellipsis if all axes are specified in addition
-    H = Axis("Height", 2)
-    W = Axis("Width", 3)
-    D = Axis("Depth", 4)
+    H, W, D = hax.make_axes(Height=2, Width=3, Depth=4)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
@@ -321,7 +300,7 @@ def test_rearrange_unused_ellipsis():
 
 
 def test_arange():
-    H = Axis("Height", 10)
+    (H,) = hax.make_axes(H=10)
 
     assert jnp.all(jnp.equal(hax.arange(H).array, jnp.arange(10)))
     assert hax.arange(H).axes == (H,)
@@ -334,8 +313,7 @@ def test_arange():
 
 
 def test_stack():
-    H = Axis("H", 4)
-    W = Axis("W", 3)
+    H, W = hax.make_axes(H=4, W=3)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W))
     named2 = hax.random.uniform(PRNGKey(1), (H, W))
@@ -350,9 +328,8 @@ def test_stack():
 
 
 def test_concatenate():
-    H1 = Axis("H", 4)
-    H2 = Axis("H", 3)
-    W = Axis("W", 3)
+    H1, W = hax.make_axes(H=4, W=3)
+    H2 = H1.resize(3)
 
     named1 = hax.random.uniform(PRNGKey(0), (H1, W))
     named2 = hax.random.uniform(PRNGKey(1), (H2, W))
@@ -392,8 +369,7 @@ def test_repeat():
     #        [3, 4],
     #        [3, 4]])
 
-    H = Axis("H", 2)
-    W = Axis("W", 2)
+    H, W = hax.make_axes(H=2, W=2)
 
     named1 = hax.named([[1, 2], [3, 4]], (H, W))
 
@@ -459,9 +435,7 @@ def test_tile():
 
 
 def test_unflatten_axis():
-    H = Axis("Height", 2)
-    W = Axis("Width", 3)
-    D = Axis("Depth", 4)
+    H, W, D = hax.make_axes(Height=2, Width=3, Depth=4)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
     flattened_HW = named1.flatten_axes((H, W), "Z")
@@ -483,9 +457,7 @@ def test_unflatten_axis():
 
 
 def test_ravel():
-    H = Axis("Height", 2)
-    W = Axis("Width", 3)
-    D = Axis("Depth", 4)
+    H, W, D = hax.make_axes(Height=2, Width=3, Depth=4)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
     raveled = named1.ravel("Z")
@@ -496,13 +468,8 @@ def test_ravel():
 
 
 def test_rename():
-    H = Axis("H", 2)
-    W = Axis("W", 3)
-    D = Axis("D", 4)
-
-    H2 = Axis("H2", 2)
-    W2 = Axis("W2", 3)
-    D2 = Axis("D2", 4)
+    H, W, D = hax.make_axes(H=2, W=3, D=4)
+    H2, W2, D2 = hax.make_axes(H2=2, W2=3, D2=4)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
@@ -517,9 +484,7 @@ def test_rename():
 
 
 def test_index():
-    H = Axis("H", 20)
-    W = Axis("W", 30)
-    D = Axis("D", 40)
+    H, W, D = hax.make_axes(H=20, W=30, D=40)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
@@ -544,9 +509,7 @@ def test_index():
 
 
 def test_index_with_tracer():
-    H = Axis("H", 20)
-    W = Axis("W", 30)
-    D = Axis("D", 40)
+    H, W, D = hax.make_axes(H=20, W=30, D=40)
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
     @jax.jit
@@ -562,12 +525,7 @@ def test_index_with_tracer():
 
 def test_index_array_slices():
     # fancier tests with array slices with named array args
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-    C = Axis("C", 40)
-    Q = Axis("Q", 50)
-    I0 = Axis("I0", 10)
+    H, W, D, C, Q, I0 = hax.make_axes(H=10, W=20, D=30, C=40, Q=50, I0=10)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D, C, Q))
     index_1 = hax.random.randint(PRNGKey(0), (I0,), 0, H.size)
@@ -586,9 +544,7 @@ def test_index_array_slices():
     # https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
     # Example
     # Let x.shape be (10, 20, 30, 40, 50) and suppose ind_1 and ind_2 can be broadcast to the shape (2, 3, 4).
-    I1 = Axis("I1", 2)
-    I2 = Axis("I2", 3)
-    I3 = Axis("I3", 4)
+    I1, I2, I3 = hax.make_axes(I1=2, I2=3, I3=4)
 
     ind_1 = hax.random.randint(PRNGKey(0), (I2, I3), 0, W.size)
     ind_2 = hax.random.randint(PRNGKey(0), (I1, I3), 0, D.size)
@@ -615,9 +571,7 @@ def test_index_array_slices():
 def test_slice_nd_shorthand_syntax():
     # syntax like arr["X", 0:10, "Y", 0:10] is supported
 
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
+    H, W, D = hax.make_axes(H=10, W=20, D=30)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
@@ -625,9 +579,7 @@ def test_slice_nd_shorthand_syntax():
 
 
 def test_slice_nd_dslice():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
+    H, W, D = hax.make_axes(H=10, W=20, D=30)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
     from haliax import ds
@@ -641,9 +593,7 @@ def test_slice_nd_dslice():
 
 def test_slice_nd_array_present_dims():
     # tests slicing with arrays that are already present in the named array, which is sometimes ok
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
+    H, W, D = hax.make_axes(H=10, W=20, D=30)
 
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
 
@@ -653,321 +603,3 @@ def test_slice_nd_array_present_dims():
     assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1.array, :, :]))
 
     # this is not ok, since the H would not be eliminated
-    with pytest.raises(ValueError):
-        named1[{W: index1}]
-
-    # this is not ok, but is trickier because the H has a different size
-    H2 = H.resize(5)
-    index2 = hax.random.randint(PRNGKey(0), (H2,), 0, H.size)
-    with pytest.raises(ValueError):
-        named1[{W: index2}]
-
-    # this is ok, since the H would be eliminated anyway
-    assert jnp.all(jnp.equal(named1[{"H": index2}].array, named1.array[index2.array, :, :]))
-
-
-def test_slice_nd_array_unnamed_slice():
-    # tests slicing with arrays that are already present in the named array, which is sometimes ok
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
-
-    index1 = jax.random.randint(PRNGKey(1), (4,), 0, H.size)
-    assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1, :, :]))
-
-    # hidden behavior: if we also pass in an H index to e.g. D, it is zipped together
-    index2 = hax.random.randint(PRNGKey(2), Axis("H", 4), 0, D.size)
-    assert jnp.all(jnp.equal(named1[{"H": index1, "D": index2}].array, named1.array[index1, :, index2.array]))
-
-    # this is different though:
-    index2r = index2.array
-    assert jnp.all(
-        jnp.equal(
-            named1[{"H": index1, "D": index2r}].array, named1.array[index1.reshape(1, -1), :, index2r.reshape(-1, 1)]
-        )
-    )
-    assert named1[{"H": index1, "D": index2r}].shape != named1[{"H": index1, "D": index2}].shape
-
-    index1 = list(index1)
-    assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1, :, :]))
-
-
-def test_full_indexing_returns_named_array():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
-    sliced = named1[{"H": 0, "W": 0, "D": 0}]
-
-    assert isinstance(sliced, NamedArray)
-    assert sliced.shape == {}
-
-
-def test_indexing_bug_from_docs():
-    X = hax.Axis("X", 10)
-    Y = hax.Axis("Y", 20)
-    Z = hax.Axis("Z", 30)
-
-    a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y, Z))
-
-    I1 = hax.Axis("I1", 5)
-    I2 = hax.Axis("I2", 5)
-    I3 = hax.Axis("I3", 5)
-    ind1 = hax.random.randint(jax.random.PRNGKey(0), (I1,), 0, 10)
-    ind2 = hax.random.randint(jax.random.PRNGKey(0), (I2, I3), 0, 20)
-
-    # assert a[{"X": ind1, "Y": ind2}].axes == (I1, I2, I3, Z)
-    assert a[{"X": ind1, "Y": ind2, "Z": 3}].axes == (I1, I2, I3)
-
-
-def test_duplicate_axis_names_in_slicing():
-    X = hax.Axis("X", 10)
-    Y = hax.Axis("Y", 20)
-    Z = hax.Axis("Z", 30)
-
-    X2 = hax.Axis("X", 5)
-    Y2 = hax.Axis("Y", 5)
-
-    a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y, Z))
-    ind1 = hax.random.randint(jax.random.PRNGKey(0), (X2,), 0, 10)
-    ind2 = hax.random.randint(jax.random.PRNGKey(0), (Y2,), 0, 10)
-
-    a[{"X": ind1, "Y": ind2}]  # returns a NamedArray with axes = Axis("X", 5), Axis("Y", 5), Axis("Z", 30)
-
-    with pytest.raises(ValueError):
-        a[{"Y": ind1}]  # error, "X" is not eliminated by the indexing operation
-
-    a[{"X": ind2, "Y": ind1}]  # ok, because X and Y are eliminated by the indexing operation
-
-
-def test_slice_old_style():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W, D), minval=0, maxval=10)
-
-    assert jnp.all(named1.slice("H", start=4, length=2).array == named1.array[4:6, :, :])
-    assert jnp.all(named1.slice("W", start=4, length=2).array == named1.array[:, 4:6, :])
-    assert jnp.all(named1.slice("D", start=4, length=2).array == named1.array[:, :, 4:6])
-
-    H2 = Axis("H2", 5)
-    W2 = Axis("W2", 10)
-    D2 = Axis("D2", 15)
-
-    assert jnp.all(named1.slice("H", H2, start=4).array == named1.array[4 : 4 + H2.size, :, :])
-    assert jnp.all(named1.slice("W", W2, start=4).array == named1.array[:, 4 : 4 + W2.size, :])
-    assert jnp.all(named1.slice("D", D2, start=4).array == named1.array[:, :, 4 : 4 + D2.size])
-
-
-def test_slice_new_style():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W, D), minval=0, maxval=10)
-
-    x1 = named1.slice({"H": 4, "W": 5, "D": 7}, length={"H": 2, "W": 3, "D": 4})
-    assert jnp.all(x1.array == named1.array[4:6, 5:8, 7:11])
-
-    with pytest.raises(TypeError):
-        named1.slice({"H": 4, "W": 5, "D": 7}, length={"H": 2, "W": 3, "D": 4}, start={"H": 1, "W": 2, "D": 3})
-
-    with pytest.raises(ValueError):
-        named1.slice({"H": 4, "W": 5, "D": 7}, length={"H": 2, "W": 3})
-
-    H2 = Axis("H2", 5)
-    W2 = Axis("W2", 10)
-    D2 = Axis("D2", 15)
-
-    x2 = named1.slice({"H": 4, "W": 5, "D": 7}, length={"H": H2, "W": W2, "D": D2})
-    assert jnp.all(x2.array == named1.array[4 : 4 + H2.size, 5 : 5 + W2.size, 7 : 7 + D2.size])
-
-
-def test_updated_slice():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    H2 = H.resize(5)
-    W2 = W.resize(10)
-    D2 = D.resize(15)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W, D), minval=0, maxval=10)
-    named2 = hax.random.randint(PRNGKey(0), (H2, W2, D2), minval=10, maxval=30)
-
-    named1_updated = named1.updated_slice({"H": 0, "W": 0, "D": 0}, named2)
-
-    assert named1_updated.axes == named1.axes
-    assert jnp.all(named1_updated["H", 0 : H2.size, "W", 0 : W2.size, "D", 0 : D2.size].array == named2.array)
-
-    # test broadcasting
-    for pair in [(H2, D2), (H2, W2), (W2, D2), (D2, H2), (D2, W2), (W2, H2)]:
-        n3 = hax.random.randint(PRNGKey(0), pair, minval=10, maxval=30)
-        named1_updated = named1.updated_slice({ax.name: 0 for ax in pair}, n3)
-        assert named1_updated.axes == named1.axes
-        assert jnp.all((named1_updated[{ax.name: slice(0, ax.size) for ax in pair}] == n3).array)
-        # check that the array outside the slice is unchanged
-        assert jnp.all(
-            (
-                named1_updated[{ax.name: slice(ax.size, None) for ax in pair}]
-                == named1[{ax.name: slice(ax.size, None) for ax in pair}]
-            ).array
-        )
-
-
-def test_updated_slice_extra_update_axis_errors():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    D = Axis("D", 30)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W, D), minval=0, maxval=10)
-    named2 = hax.random.randint(PRNGKey(0), (H, W, D), minval=10, maxval=30)
-
-    with pytest.raises(ValueError):
-        named1.updated_slice({"H": 0, "W": 0, "D": 0, "extra": 0}, named2)
-
-    with pytest.raises(ValueError):
-        named3 = hax.random.randint(PRNGKey(0), (H, W), minval=10, maxval=30)
-        named3.updated_slice({"H": 0, "W": 0}, named2)
-
-
-def test_order_of_transpose_add():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W), minval=0, maxval=10)
-    named2 = hax.random.randint(PRNGKey(0), (W, H), minval=10, maxval=30)
-
-    assert (named1 + named2).axes == (H, W)
-    assert jnp.all((named1 + named2).array == named1.array + named2.array.T)
-
-
-def test_nice_short_string_in_named_array():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W), minval=0, maxval=10)
-
-    assert str(named1).startswith("NamedArray(int32{'H': 10, 'W': 20}")
-
-
-def test_nice_short_string_in_named_array_in_eqx_module():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-
-    named1 = hax.random.randint(PRNGKey(0), (H, W), minval=0, maxval=10)
-
-    class TestModule(eqx.Module):
-        named1: NamedArray
-
-    mod = TestModule(named1)
-
-    assert str(mod).startswith("TestModule(named1=Named(int32{'H': 10, 'W': 20}))")
-
-
-def test_named_arrays_work_in_eqxi_while_loop():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W))
-
-    import equinox.internal as eqxi
-
-    def body_fun(t):
-        i, named1 = t
-        return i + 1, named1 + named1
-
-    def cond_fun(t):
-        i, named1 = t
-        return i < 10
-
-    def loss_fun(named1):
-        i, named1 = eqxi.while_loop(cond_fun, body_fun, (0, named1), kind="checkpointed", max_steps=10)
-        return named1.sum().scalar()
-
-    grad_fun = eqx.filter_value_and_grad(loss_fun)
-
-    grad_fun(named1)
-
-
-def test_at_for_in_placeish():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W))
-
-    named1_at = named1.at[H, 0].set(0)
-
-    assert jnp.all(jnp.equal(named1_at[H, 0].array, 0))
-    assert jnp.all(named1_at[H, 1:].array == named1[H, 1:].array)
-
-    # test add, multiply, power, etc.
-    named1_at = named1.at[H, 0].add(1)
-    assert jnp.all(named1_at.array == named1.array.at[0].add(1))
-
-    named1_at = named1.at[H, 0].multiply(2)
-    assert jnp.all(named1_at.array == named1.array.at[0].multiply(2))
-
-    named1_at = named1.at[H, 0].power(2)
-    assert jnp.all(named1_at.array == named1.array.at[0].power(2))
-
-    named1_at = named1.at[H, 0].divide(2)
-    assert jnp.all(named1_at.array == named1.array.at[0].divide(2))
-
-    named1_at = named1.at[H, 0].apply(hax.square)
-    assert jnp.all(named1_at.array == named1.array.at[0].apply(jnp.square))
-
-    named1_at = named1.at[H, 0].max(0.5)
-    assert jnp.all(named1_at.array == named1.array.at[0].max(0.5))
-
-    named1_at = named1.at[H, 0].min(0.5)
-    assert jnp.all(named1_at.array == named1.array.at[0].min(0.5))
-
-
-def test_at_with_fancy_indexing():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    I0 = Axis("I0", 5)
-    I1 = Axis("I1", 5)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W))
-    ind1 = hax.random.randint(PRNGKey(0), (I0,), 0, H.size)
-    ind2 = hax.random.randint(PRNGKey(0), (I1,), 0, W.size)
-
-    named1_at = named1.at[H, ind1].set(0)
-    assert jnp.all(named1_at.array == named1.array.at[ind1.array].set(0))
-
-    named1_at = named1.at[H, ind1].add(1, mode="clip")
-    assert jnp.all(named1_at.array == named1.array.at[ind1.array].add(1, mode="clip"))
-
-    named1_at = named1.at[H, ind1, W, ind2].set(0)
-    assert jnp.all(named1_at.array == named1.array.at[ind1.array.reshape(-1, 1), ind2.array.reshape(1, -1)].set(0))
-
-    # dslices
-    from haliax import ds
-
-    named1_at = named1.at[H, ds(3, 5)].set(0)
-    assert jnp.all(named1_at.array == named1.array.at[3:8].set(0))
-
-    named1_at = named1.at[H, ds(3, 5), W, ind2].power(2)
-    assert jnp.all(named1_at.array == named1.array.at[3:8, ind2.array].power(2))
-
-
-def test_slice_dslice_and_array():
-    H = Axis("H", 10)
-    W = Axis("W", 20)
-    I0 = Axis("I0", 5)
-
-    named1 = hax.random.uniform(PRNGKey(0), (H, W))
-    ind2 = hax.random.randint(PRNGKey(0), (I0,), 0, W.size)
-
-    from haliax import ds
-
-    named1.array.at[3:8, ind2.array].add(jnp.full((5, 5), 2))
-
-    named1_at = named1.at[H, ds(3, 5), W, ind2].add(2)
-    assert jnp.all(named1_at.array == named1.array.at[3:8, ind2.array].add(2))

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_attention.py

@@ -93,8 +93,7 @@ def test_alibi_attention_compared_to_hf():
     import torch
     from transformers.models.bloom.modeling_bloom import build_alibi_tensor
 
-    L = hax.Axis("L", 128)
-    H = hax.Axis("NumHeads", 16)
+    L, H = hax.make_axes(L=1, H=16)
 
     # Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)
     torch_tensor = (
@@ -107,7 +106,7 @@ def test_alibi_attention_compared_to_hf():
 
 
 def test_fcm_attention_mask():
-    KeyPos = hax.Axis("KeyPos", 20)
+    KeyPos, QueryPos, Head = hax.make_axes(KeyPos=20, QueryPos=10, Head=8)
 
     mask = forgetful_causal_mask(KeyPos, mask_prob=0.6, sample_prob=False, key=PRNGKey(0))
 
@@ -116,9 +115,6 @@ def test_fcm_attention_mask():
 
     assert mask.astype(float).sum().item() <= KeyPos.size
 
-    QueryPos = hax.Axis("QueryPos", 10)
-    Head = hax.Axis("Head", 8)
-
     query = hax.arange(QueryPos).broadcast_axis(Head)
     key = hax.arange(KeyPos).broadcast_axis(Head)
 

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_axis.py

@@ -1,12 +1,10 @@
 import pytest
 
-from haliax.axis import Axis, eliminate_axes, rearrange_for_partial_order
+from haliax.axis import Axis, eliminate_axes, make_axes, rearrange_for_partial_order
 
 
 def test_eliminate_axes():
-    H = Axis("H", 3)
-    W = Axis("W", 4)
-    C = Axis("C", 5)
+    H, W, C = make_axes(H=3, W=4, C=5)
 
     assert eliminate_axes((H, W), (H,)) == (W,)
     assert eliminate_axes((H, W), (W,)) == (H,)

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_einsum.py

@@ -273,12 +273,7 @@ def test_einsum_various_errors():
 
 
 def test_einsum_examples():
-
-    Batch = hax.Axis("batch", 32)
-    Embed = hax.Axis("embed", 64)
-    H = hax.Axis("h", 16)
-    W = hax.Axis("w", 16)
-    C = hax.Axis("c", 3)
+    Batch, Embed, H, W, C = hax.make_axes(batch=32, embed=64, h=16, w=16, c=3)
 
     # for jax
     im = jnp.zeros((32, 16, 16, 3))

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_nn.py

@@ -46,8 +46,7 @@ def test_dropout():
 
 
 def test_one_hot():
-    i = Axis("i", 3)
-    c = Axis("c", 3)
+    i, c = hax.make_axes(i=3, c=3)
     actual = hax.nn.one_hot(hax.NamedArray(jnp.array([0, 1, 2]), (i,)), c)
     expected = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
 
@@ -61,16 +60,14 @@ def test_one_hot():
 
 
 def test_one_hot_out_of_bound():
-    i = Axis("i", 2)
-    c = Axis("c", 3)
+    i, c = hax.make_axes(i=2, c=3)
     actual = hax.nn.one_hot(hax.NamedArray(jnp.array([-1, 3]), (i,)), c)
     expected = jnp.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
     assert jnp.all(jnp.isclose(actual.array, expected))
 
 
 def test_standardize():
-    b = Axis("b", 2)
-    c = Axis("c", 3)
+    b, c = hax.make_axes(b=2, c=3)
     actual = hax.nn.standardize(hax.NamedArray(jnp.array([0, 1, 2]), (c,)), c)
     expected = jax.nn.standardize(jnp.array([0, 1, 2]), axis=0)
 
@@ -113,11 +110,7 @@ def test_standardize():
 @pytest.mark.parametrize("depth", [0, 1, 2, 3, 4, 5])
 def test_mlp(depth):
     key = jrandom.PRNGKey(0)
-    H = Axis("H", 10)
-    C = Axis("C", 12)
-    W = Axis("W", 14)
-
-    E = Axis("E", 16)
+    H, C, W, E = hax.make_axes(H=10, C=12, W=14, E=16)
 
     hax_mlp = hax.nn.MLP.init((H, C, W), E, width=8, depth=depth, key=key)
     x = hax.random.uniform(key, (H, C, W))
@@ -150,10 +143,7 @@ def test_mlp(depth):
 
 
 def test_linear_has_no_function_leaves_by_default():
-    H = Axis("H", 10)
-    C = Axis("C", 12)
-    W = Axis("W", 14)
-    E = Axis("E", 16)
+    H, C, W, E = hax.make_axes(H=10, C=12, W=14, E=16)
 
     hax_linear = hax.nn.Linear.init((H, C, W), E, key=jrandom.PRNGKey(0))
     assert all(not isinstance(v, Callable) for v in jax.tree_util.tree_leaves(hax_linear))  # type: ignore

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_rearrange.py

@@ -7,11 +7,7 @@ from haliax._src.rearrange import einops_rearrange
 
 
 # some axes
-W = Axis("W", 4)
-H = Axis("H", 6)
-C = Axis("C", 3)
-D = Axis("D", 2)
-B = Axis("B", 5)
+H, W, C, D, B = hax.make_axes(H=6, W=4, C=3, D=2, B=5)
 Q = hax.Axis("Q", B.size * H.size)
 E = Axis("E", H.size * W.size * D.size)
 

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_specialized_fns.py

@@ -1,14 +1,13 @@
 import jax
 import jax.numpy as jnp
 
+import haliax
 import haliax.specialized_fns as hfns
-from haliax import Axis, NamedArray
+from haliax import NamedArray
 
 
 def test_top_k():
-    H = Axis("H", 5)
-    W = Axis("W", 6)
-    D = Axis("D", 7)
+    H, W, D = haliax.make_axes(H=3, W=4, D=5)
 
     rand = jax.random.uniform(jax.random.PRNGKey(0), (H.size, W.size, D.size))
     n_rand = NamedArray(rand, (H, W, D))

{haliax-1.4.dev315 → haliax-1.4.dev317}/tests/test_tree_util.py

@@ -11,10 +11,7 @@ from haliax import Axis
 
 
 def test_resize_axis():
-
-    A = hax.Axis("A", 10)
-    B = hax.Axis("B", 20)
-    C = hax.Axis("C", 30)
+    A, B, C = hax.make_axes(A=10, B=20, C=30)
 
     class Module(eqx.Module):
         name1: hax.NamedArray

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

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