haliax 1.4.dev354__tar.gz → 1.4.dev356__tar.gz

{haliax-1.4.dev354 → haliax-1.4.dev356}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev354
+Version: 1.4.dev356
 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.dev354 → haliax-1.4.dev356}/docs/indexing.md

@@ -287,3 +287,42 @@ operation more effectively.)
 
     It's worth emphasizing that these functions are typically compiled to scatter-add and friends (as appropriate).
     This is the preferred way to do scatter/gather operations in JAX, as well as in Haliax.
+
+## Scatter/Gather
+
+Haliax supports scatter/gather semantics in its indexing operations. When an axis
+is indexed by another NamedArray (or a 1-D JAX array), the values of that axis
+are gathered according to the index array and the axes of the indexer are
+inserted into the result.
+
+```python
+import haliax as hax
+import jax.numpy as jnp
+
+B, S, V = Axis("batch", 4), Axis("seq", 3), Axis("vocab", 7)
+x = hax.arange((B, S, V))
+idx = hax.arange((B, S), dtype=jnp.int32) % V.size
+
+out = x["vocab", idx]
+```
+
+Here `out` has axes `(B, S)` and its values match `jax.numpy.take_along_axis`
+on the underlying ndarray.
+
+For scatter-style updates where each batch writes to a different position, use
+[`updated_slice`][haliax.updated_slice]:
+
+```python
+Batch = hax.Axis("batch", 2)
+Seq = hax.Axis("seq", 5)
+New = hax.Axis("seq", 2)
+
+cache = hax.zeros((Batch, Seq), dtype=int)
+lengths = hax.named([1, 3], axis=Batch)
+kv = hax.named([[1, 2], [3, 4]], axis=(Batch, New))
+
+result = updated_slice(cache, {"seq": lengths}, kv)
+```
+
+This inserts `[1, 2]` starting at position `1` in batch `0` and `[3, 4]` starting
+at position `3` in batch `1`.

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

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

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

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