haliax 1.4.dev355__tar.gz → 1.4.dev356__tar.gz

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

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

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

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

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