haliax 1.4.dev355__tar.gz → 1.4.dev359__tar.gz

{haliax-1.4.dev355 → haliax-1.4.dev359}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev355
+Version: 1.4.dev359
 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.dev359/src/haliax/__about__.py

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

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/__init__.py

@@ -212,7 +212,7 @@ def repeat(
     a: NamedArray, repeats: int | jnp.ndarray, axis: AxisSelector, total_repeat_length: Optional[int] = None
 ) -> NamedArray:
     """Version of [jax.numpy.repeat][] that returns a NamedArray"""
-    index = a._lookup_indices(axis)
+    index = a.axis_indices(axis)
     if index is None:
         raise ValueError(f"Axis {axis} not found in array {a}")
 
@@ -235,7 +235,7 @@ def tile(a: NamedArray, reps: dict[AxisSelector, int]) -> NamedArray:
     new_dims = []
     dim_reps = [1] * len(a.axes)
     for ax, i in reps.items():
-        index = a._lookup_indices(ax)
+        index = a.axis_indices(ax)
         if index is None:
             new_dims.append(Axis(axis_name(ax), i))
         else:
@@ -265,7 +265,7 @@ def concatenate(axis: AxisSelector, arrays: Sequence[NamedArray]) -> NamedArray:
     if len(arrays) == 0:
         return zeros(axis)
 
-    axis_index = arrays[0]._lookup_indices(aname)
+    axis_index = arrays[0].axis_indices(aname)
     if axis_index is None:
         raise ValueError(f"Axis {aname} not found in 0th array {arrays[0]}")
 

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/core.py

@@ -1,6 +1,7 @@
 import contextlib
 import functools as ft
 import typing
+import warnings
 from dataclasses import dataclass
 from math import prod
 from types import EllipsisType
@@ -142,7 +143,7 @@ class NamedArray:
 
     def has_axis(self, axis: AxisSelection) -> bool:
         """Returns true if the given axis is present in this NamedArray."""
-        return self._lookup_indices(axis) is not None
+        return self.axis_indices(axis) is not None
 
     @overload
     def axis_size(self, axis: AxisSelector) -> int:  # type: ignore
@@ -156,7 +157,7 @@ class NamedArray:
         """
         Returns the size of the given axis, or a tuple of sizes if given multiple axes.
         """
-        indices = self._lookup_indices(axis)
+        indices = self.axis_indices(axis)
         if isinstance(indices, int):
             return self.axes[indices].size
         elif indices is None:
@@ -184,7 +185,7 @@ class NamedArray:
 
         Raises a ValueError if any of the axes are not found.
         """
-        indices = self._lookup_indices(axes)
+        indices = self.axis_indices(axes)
         if isinstance(indices, int):
             return self.axes[indices]
         elif indices is None:
@@ -235,6 +236,27 @@ class NamedArray:
         For a single axis, returns an int corresponding to the index of the axis.
         For multiple axes, returns a tuple of ints corresponding to the indices of the axes.
 
+        If the axis is not present, returns None for that position
+        """
+        warnings.warn(
+            "_lookup_indices() is deprecated, use axis_indices() instead",
+            DeprecationWarning,
+        )
+        return self.axis_indices(axis)
+
+    @overload
+    def axis_indices(self, axis: AxisSelector) -> Optional[int]:  # type: ignore
+        ...
+
+    @overload
+    def axis_indices(self, axis: Sequence[AxisSelector]) -> Tuple[Optional[int], ...]:
+        ...
+
+    def axis_indices(self, axis: AxisSelection) -> Union[Optional[int], Tuple[Optional[int], ...]]:
+        """
+        For a single axis, returns an int corresponding to the index of the axis.
+        For multiple axes, returns a tuple of ints corresponding to the indices of the axes.
+
         If the axis is not present, returns None for that position
         """
         if isinstance(axis, Axis):
@@ -262,7 +284,7 @@ class NamedArray:
             except ValueError:
                 return None
         else:
-            return tuple(self._lookup_indices(a) for a in axis)
+            return tuple(self.axis_indices(a) for a in axis)
 
     # Axis rearrangement
     @typing.overload
@@ -675,7 +697,7 @@ def take(array: NamedArray, axis: AxisSelector, index: Union[int, NamedArray]) -
 
     if index is a NamedArray, then those axes are added to the output array
     """
-    axis_index = array._lookup_indices(axis)
+    axis_index = array.axis_indices(axis)
     if axis_index is None:
         raise ValueError(f"axis {axis} not found in {array}")
 
@@ -696,14 +718,14 @@ def take(array: NamedArray, axis: AxisSelector, index: Union[int, NamedArray]) -
 
         if intersecting_axes:
             # if the eliminated axis is also in the index, we rename it to a dummy axis that we can broadcast over it
-            need_to_use_dummy_axis = index._lookup_indices(axis.name) is not None
+            need_to_use_dummy_axis = index.axis_indices(axis.name) is not None
             if need_to_use_dummy_axis:
                 index = index.rename({axis.name: "__DUMMY_" + axis.name})
             array = haliax.broadcast_to(array, index.axes, ensure_order=False, enforce_no_extra_axes=False)
             new_axes = eliminate_axes(array.axes, axis)
             index = haliax.broadcast_to(index, new_axes, ensure_order=True, enforce_no_extra_axes=True)
 
-            axis_index = array._lookup_indices(axis)  # if it moved
+            axis_index = array.axis_indices(axis)  # if it moved
             index_array = jnp.expand_dims(index.array, axis=axis_index)
             new_array = jnp.take_along_axis(array.array, index_array, axis=axis_index)
             new_array = jnp.squeeze(new_array, axis=axis_index)
@@ -792,7 +814,7 @@ def _slice_old(
     Note:
         This method is basically a wrapper around jax.lax.dynamic_slice_in_dim.
     """
-    axis_index = array._lookup_indices(axis)
+    axis_index = array.axis_indices(axis)
     if axis_index is None:
         raise ValueError(f"axis {axis} not found in {array}")
 
@@ -824,7 +846,7 @@ def _slice_new(
     new_lengths = [axis.size for axis in array.axes]
 
     for axis, s in start.items():
-        axis_index = array._lookup_indices(axis_name(axis))
+        axis_index = array.axis_indices(axis_name(axis))
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
 
@@ -883,14 +905,14 @@ def updated_slice(
         f = updated_slice
         for axis_name in map_axes:
             # make sure that axis_name is in `array`. otherwise it doesn't make sense to vmap over it
-            if array._lookup_indices(axis_name) is None:
+            if array.axis_indices(axis_name) is None:
                 raise ValueError(f"axis {axis_name} not found in original array's axes: {array.shape}")
             f = haliax.vmap(f, axis=axis_name)
         return f(array, start, update)
 
     array_slice_indices = [0] * len(array.axes)
     for axis, s in start.items():
-        axis_index = array._lookup_indices(haliax.axis_name(axis))
+        axis_index = array.axis_indices(haliax.axis_name(axis))
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
         if isinstance(s, NamedArray):  # this can happen in the vmap case
@@ -900,7 +922,7 @@ def updated_slice(
 
         array_slice_indices[axis_index] = s
         total_length = array.axes[axis_index].size
-        update_axis = update._lookup_indices(haliax.axis_name(axis))
+        update_axis = update.axis_indices(haliax.axis_name(axis))
 
         # if s is a tracer we can't check the size
         if update_axis is None:
@@ -923,7 +945,7 @@ def updated_slice(
             )
         broadcasted_axes = []
         for ax in array.axes:
-            upd_ax = update._lookup_indices(ax.name)
+            upd_ax = update.axis_indices(ax.name)
             broadcasted_axes.append(ax if upd_ax is None else update.axes[upd_ax])
         update = haliax.broadcast_to(update, broadcasted_axes, enforce_no_extra_axes=True)
         upd_arr = update.array
@@ -967,7 +989,7 @@ def _compute_new_axes_and_slices_for_index(
     index_axis_names = set()
 
     for axis, slice_ in slices.items():
-        axis_index = array._lookup_indices(axis)
+        axis_index = array.axis_indices(axis)
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
         if isinstance(slice_, py_slice) or isinstance(slice_, dslice) or is_pallas_dslice(slice_):
@@ -1111,7 +1133,7 @@ def split(a: NamedArray, axis: AxisSelector, new_axes: Sequence[Axis]) -> Sequen
         new_axes (Sequence[Axis]): the axes to split into. Must have the same total length as the axis being split.
     """
     # check the lengths of the new axes
-    index = a._lookup_indices(axis)
+    index = a.axis_indices(axis)
     if index is None:
         raise ValueError(f"Axis {axis} not found in {a.axes}")
 
@@ -1136,7 +1158,7 @@ def unbind(array: NamedArray, axis: AxisSelector) -> List[NamedArray]:
     Unbind an array along an axis, returning a list of NamedArrays, one for each position on that axis.
     Analogous to torch.unbind or np.rollaxis
     """
-    axis_index = array._lookup_indices(axis)
+    axis_index = array.axis_indices(axis)
     if axis_index is None:
         raise ValueError(f"axis {axis} not found in {array}")
     new_axes = array.axes[:axis_index] + array.axes[axis_index + 1 :]
@@ -1153,7 +1175,7 @@ def roll(array: NamedArray, shift: Union[int, Tuple[int, ...]], axis: AxisSelect
     """
     Roll an array along an axis or axes. Analogous to np.roll
     """
-    axis_indices = array._lookup_indices(axis)
+    axis_indices = array.axis_indices(axis)
     if axis_indices is None:
         raise ValueError(f"axis {axis} not found in {array}")
     return NamedArray(jnp.roll(array.array, shift, axis_indices), array.axes)
@@ -1331,7 +1353,7 @@ def unflatten_axis(array: NamedArray, axis: AxisSelector, new_axes: AxisSpec) ->
     """
     Split an axis into a sequence of axes. The old axis must have the same size as the product of the new axes.
     """
-    old_index = array._lookup_indices(axis)
+    old_index = array.axis_indices(axis)
     if old_index is None:
         raise ValueError(f"Axis {axis} not found in {array}")
 

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/hof.py

@@ -61,8 +61,10 @@ def vmap(
         for ax in reversed(axes):
             mapped = vmap(mapped, ax, default=default, args=args, kwargs=kwargs)
         return mapped
-    else:
+    elif len(axes) == 1:  # type: ignore
         axis = axes[0]
+    else:
+        return fn
 
     signature = inspect.signature(fn)
 
@@ -83,7 +85,7 @@ def vmap(
 
     def _index_of_batch_axis(array, default):
         if isinstance(array, NamedArray):
-            return array._lookup_indices(axis)
+            return array.axis_indices(axis)
         elif callable(default):
             return default(array)
         else:

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/nn/__init__.py

@@ -23,6 +23,7 @@ from .activations import (
     quick_gelu,
     relu,
     relu6,
+    relu_squared,
     selu,
     sigmoid,
     silu,
@@ -94,6 +95,7 @@ __all__ = [
     "quick_gelu",
     "glu",
     "relu6",
+    "relu_squared",
     "sigmoid",
     "soft_sign",
     "softplus",

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/nn/activations.py

@@ -87,3 +87,13 @@ def glu(x: NamedArray, axis: Axis) -> NamedArray:
 
 def quick_gelu(x):
     return x * sigmoid(1.702 * x)
+
+
+
+def relu_squared(x: A) -> A:
+    """ReLU squared activation function. jnp.square(jnp.maximum(0, x))"""
+
+    def _fn(a):
+        return jnp.square(jnn.relu(a))
+
+    return typing.cast(A, wrap_elemwise_unary(_fn, x))

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/nn/attention.py

@@ -111,7 +111,7 @@ def dot_product_attention(
     KPos = ensure_tuple(key.resolve_axis(KPos))
     # any axis in KPos that's in query is a problem
     for axis in KPos:
-        if axis in query.axes:
+        if query.has_axis(axis):
             raise ValueError(
                 f"Axis {axis} in KPos is also in query. Attended-to axes must be distinct from query axis"
             )

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/nn/normalization.py

@@ -163,7 +163,7 @@ def standardize(
     """Analogous to [jax.nn.standardize][], but with support for NamedArrays."""
     x, mean, variance, where = haliax.broadcast_arrays(x, mean, variance, where)  # type: ignore
     raw_x, mean, variance, where = unwrap_namedarrays(x, mean, variance, where)
-    axis_indices = x._lookup_indices(axis)
+    axis_indices = x.axis_indices(axis)
 
     plain = jnn.standardize(raw_x, axis_indices, mean=mean, variance=variance, epsilon=epsilon, where=where)
     return NamedArray(plain, x.axes)

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/ops.py

@@ -11,8 +11,8 @@ from .jax_utils import is_scalarish
 
 def trace(array: NamedArray, axis1: AxisSelector, axis2: AxisSelector, offset=0, dtype=None) -> NamedArray:
     """Compute the trace of an array along two named axes."""
-    a1_index = array._lookup_indices(axis1)
-    a2_index = array._lookup_indices(axis2)
+    a1_index = array.axis_indices(axis1)
+    a2_index = array.axis_indices(axis2)
 
     if a1_index is None:
         raise ValueError(f"Axis {axis1} not found in array. Available axes: {array.axes}")
@@ -139,7 +139,7 @@ def pad_left(array: NamedArray, axis: Axis, new_axis: Axis, value=0) -> NamedArr
     if amount_to_pad_to < 0:
         raise ValueError(f"Cannot pad {axis} to {new_axis}")
 
-    idx = array._lookup_indices(axis)
+    idx = array.axis_indices(axis)
 
     padding = [(0, 0)] * array.ndim
     if idx is None:

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/random.py

@@ -218,7 +218,7 @@ def choice(
     :return: Array with shape `shape` + (`a.axes` - `axis`)
     """
 
-    index = a._lookup_indices(axis)
+    index = a.axis_indices(axis)
     assert index is not None, f"axis {axis} not in a"
 
     shape = ensure_tuple(shape)
@@ -261,7 +261,7 @@ def categorical(key, logits: NamedArray, axis: AxisSelector, shape: Optional[Axi
 
     logits = logits.broadcast_axis(shape)
 
-    index = logits._lookup_indices(axis)
+    index = logits.axis_indices(axis)
     assert index is not None, f"axis {axis} not in logits"
 
     jax_shape = to_jax_shape(shape)
@@ -280,7 +280,7 @@ def gumbel(key, shape: AxisSpec, dtype=float):
 
 @named_call
 def permutation(key, x: NamedArray, axis: AxisSelector, independent: bool = False):
-    axis_index = x._lookup_indices(axis)
+    axis_index = x.axis_indices(axis)
     jax_array = jrandom.permutation(key, x.array, axis_index, independent=independent)
     return haliax.auto_sharded(NamedArray(jax_array, x.axes))
 

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/specialized_fns.py

@@ -22,7 +22,7 @@ def top_k(
         NamedArray: array with the top k elements along the given axis
         NamedArray: array with the top k elements' indices along the given axis
     """
-    pos = arr._lookup_indices(axis)
+    pos = arr.axis_indices(axis)
     if pos is None:
         raise ValueError(f"Axis {axis} not found in {arr}")
     new_array = jnp.moveaxis(arr.array, pos, -1)  # move axis to the last position

{haliax-1.4.dev355 → haliax-1.4.dev359}/src/haliax/wrap.py

@@ -52,7 +52,7 @@ def wrap_reduction_call(
                 axis = ensure_tuple(axis)
                 if single_axis_only and len(axis) > 1:
                     raise ValueError(f"{fn.__name__} only supports a single axis")
-                indices = a._lookup_indices(axis)
+                indices = a.axis_indices(axis)
                 if indices is None or any(x is None for x in indices):
                     raise ValueError(f"axis {axis} is not in {a.axes}")
                 new_axes = [ax for ax in a.axes if not selects_axis(axis, ax)]
@@ -74,7 +74,7 @@ def wrap_axiswise_call(fn, a, axis: Optional[AxisSelection], *, single_axis_only
         if axis is None:
             return fn(a.array, axis=None, **kwargs)
         else:
-            indices = ensure_tuple(a._lookup_indices(axis))
+            indices = ensure_tuple(a.axis_indices(axis))
             if any(x is None for x in indices):
                 raise ValueError(f"axis {axis} is not in {a.axes}")
             if len(indices) == 1:

{haliax-1.4.dev355 → haliax-1.4.dev359}/tests/test_nn.py

@@ -147,3 +147,55 @@ def test_linear_has_no_function_leaves_by_default():
 
     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
+
+
+@pytest.mark.parametrize(
+    "input_data, axes",
+    [
+        (jnp.array([-2.0, -1.0, 0.0, 1.0, 2.0]), (hax.Axis("X", 5),)),
+        (jnp.array([[1.0, -1.0], [0.0, 2.0]]), (hax.Axis("Y", 2), hax.Axis("Z", 2))),
+        (jnp.array([jnp.nan, 1.0, -1.0]), (hax.Axis("A", 3),)),
+        (jnp.array([jnp.inf, -jnp.inf, 0.0]), (hax.Axis("B", 3),)),
+    ],
+)
+@pytest.mark.parametrize("dtype", [jnp.float16, jnp.float32, jnp.bfloat16])
+@pytest.mark.parametrize("use_jit", [False, True])
+def test_relu_squared_robust(input_data, axes, dtype, use_jit):
+    input_data = input_data.astype(dtype)
+    x = hax.named(input_data, axes)
+
+    # Manually compute the expected output using the base JAX functions
+    expected_raw = jnp.square(jax.nn.relu(input_data))
+    expected = hax.named(expected_raw, axes)
+
+    f = hax.nn.relu_squared
+    if use_jit:
+        f = hax.named_jit(f)
+
+    # Apply the relu_squared function
+    actual = f(x)
+
+    # Check that the output is a NamedArray with the correct axes and dtype
+    assert isinstance(actual, hax.NamedArray)
+    assert actual.axes == expected.axes
+    assert actual.dtype == expected.dtype
+
+    # Check that the values are correct, handling NaNs correctly
+    assert jnp.allclose(actual.array, expected.array, equal_nan=True)
+
+
+@pytest.mark.parametrize("use_jit", [False, True])
+def test_relu_squared_scalar(use_jit):
+    f = hax.nn.relu_squared
+    if use_jit:
+        f = jax.jit(f)
+
+    x = 5.0
+    expected = 25.0
+    actual = f(x)
+    assert jnp.allclose(actual, expected)
+
+    x_neg = -5.0
+    expected_neg = 0.0
+    actual_neg = f(x_neg)
+    assert jnp.allclose(actual_neg, expected_neg)

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

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