haliax 1.4.dev290__tar.gz → 1.4.dev292__tar.gz

{haliax-1.4.dev290 → haliax-1.4.dev292}/PKG-INFO

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

@@ -153,9 +153,11 @@ It's not a huge improvement, but it's a bit more convenient.
 ## Advanced Indexing
 
 NumPy's [Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing) is supported, though we use named arrays for the indices instead of normal arrays.
-In NumPy, the indexed arrays much be broadcastable to the same shape.Advanced indexing in Haliax is similar, except that
-they follow Haliax's broadcasting rules, meaning that the axis names determine broadcasting.Axes with the same name must
-have the same size.
+(Though, as noted above, you can use 1-D JAX arrays of integers as well.)
+In NumPy, the indexed arrays must be broadcastable to the same shape. Advanced indexing in Haliax is similar,
+except that it follows Haliax's broadcasting rules, meaning that shared names are broadcasted together,
+while non-shared names are treated as separate axes and are cross-producted.
+In particular, axes with the same name must have the same size.
 
 ```python
 import haliax as hax

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

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

{haliax-1.4.dev290 → haliax-1.4.dev292}/src/haliax/ops.py

@@ -5,7 +5,7 @@ import jax
 import jax.numpy as jnp
 
 from .axis import Axis, AxisSelector
-from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes
+from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
 from .jax_utils import is_scalarish
 
 
@@ -70,9 +70,21 @@ def where(
             for idx in jnp.where(condition.array, size=new_axis.size, fill_value=fill_value)
         )
 
+    # if x or y is a NamedArray, the other must be as well. wrap as needed for scalars
+
     if is_scalarish(condition):
-        if x is None:
+        if x is None or y is None:
             raise ValueError("Must specify x and y when condition is a scalar")
+
+        if isinstance(x, NamedArray) and not isinstance(y, NamedArray):
+            if not is_scalarish(y):
+                raise ValueError("y must be a NamedArray or scalar if x is a NamedArray")
+            y = named(y, ())
+        elif isinstance(y, NamedArray) and not isinstance(x, NamedArray):
+            if not is_scalarish(x):
+                raise ValueError("x must be a NamedArray or scalar if y is a NamedArray")
+            x = named(x, ())
+        x, y = broadcast_arrays(x, y)
         return jax.lax.cond(condition, lambda _: x, lambda _: y, None)
 
     condition, x, y = broadcast_arrays(condition, x, y)  # type: ignore

{haliax-1.4.dev290 → haliax-1.4.dev292}/tests/test_ops.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 import jax.numpy as jnp
 import pytest
 from jax.random import PRNGKey
@@ -105,19 +107,27 @@ def test_add_no_overlap():
 # TODO: tests for other ops:
 
 
-def test_where():
+@pytest.mark.parametrize("use_jit", [False, True])
+def test_where(use_jit):
     Height = Axis("Height", 10)
     Width = Axis("Width", 3)
     Depth = Axis("Depth", 4)
 
+    hax_where: Callable = hax.where
+    if use_jit:
+        hax_where = hax.named_jit(hax_where)
+
     named1 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
     named2 = hax.random.uniform(PRNGKey(1), (Height, Width, Depth))
 
-    named3 = hax.where(named1 > named2, named1, named2)
+    hax_where(0.0, named1, 0.0)
+
+    named3 = hax_where(named1 > named2, named1, named2)
+
     assert jnp.all(jnp.isclose(named3.array, jnp.where(named1.array > named2.array, named1.array, named2.array)))
 
     named2_reorder = named2.rearrange((Width, Height, Depth))
-    named4 = hax.where(named1 > named2_reorder, named1, named2_reorder)
+    named4 = hax_where(named1 > named2_reorder, named1, named2_reorder)
     named4 = named4.rearrange((Height, Width, Depth))
     assert jnp.all(jnp.isclose(named4.array, jnp.where(named1.array > named2.array, named1.array, named2.array)))
 
@@ -125,7 +135,7 @@ def test_where():
     named5 = hax.random.uniform(PRNGKey(1), (Height, Width))
     named6 = hax.random.uniform(PRNGKey(2), Width)
 
-    named7 = hax.where(named5 > named6, named5, named6)
+    named7 = hax_where(named5 > named6, named5, named6)
     named7 = named7.rearrange((Height, Width))
     assert jnp.all(jnp.isclose(named7.array, jnp.where(named5.array > named6.array, named5.array, named6.array)))
 
@@ -134,12 +144,12 @@ def test_where():
     named6 = hax.random.uniform(PRNGKey(2), (Width, Depth))
 
     with pytest.raises(ValueError):
-        _ = hax.where(named5 > named6, named5, named6)
+        _ = hax_where(named5 > named6, named5, named6)
 
     # now single argument mode
     Volume = hax.Axis("Volume", Height.size * Width.size * Depth.size)
     named7 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
-    named8, named9, named10 = hax.where(named7 > 0.5, fill_value=-1, new_axis=Volume)
+    named8, named9, named10 = hax_where(named7 > 0.5, fill_value=-1, new_axis=Volume)
     assert jnp.all((named7[{"Height": named8, "Width": named9, "Depth": named10}] > 0.5).array)
 
 

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

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