haliax 1.4.dev313__tar.gz → 1.4.dev314__tar.gz

{haliax-1.4.dev313 → haliax-1.4.dev314}/PKG-INFO

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

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

{haliax-1.4.dev313 → haliax-1.4.dev314}/src/haliax/core.py

@@ -952,6 +952,7 @@ def _compute_new_axes_and_slices_for_index(
             # we allow this if it's a 0-d or 1-d array
             if slice_.ndim == 0:
                 ordered_slices[axis_index] = slice_
+                kept_axes[axis_index] = False
             elif slice_.ndim == 1:
                 # we allow this if it's a 1-d array, in which case we treat it as sugar for NamedArray(slice_, sliced-axis)
                 ordered_slices[axis_index] = haliax.named(slice_, axis_name(axis))

{haliax-1.4.dev313 → haliax-1.4.dev314}/tests/core_test.py

@@ -543,6 +543,23 @@ def test_index():
     assert jnp.all(jnp.equal(named1[{"H": 0, "W": 0, "D": 0}], named1.array[0, 0, 0]))
 
 
+def test_index_with_tracer():
+    H = Axis("H", 20)
+    W = Axis("W", 30)
+    D = Axis("D", 40)
+    named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
+
+    @jax.jit
+    def f(idx):
+        return named1["H", idx]
+
+    idx = jnp.array([1, 2, 3])
+    assert jnp.all(jnp.equal(f(idx).array, named1.array[1:4, :, :]))
+
+    idx = jnp.array(0)
+    assert jnp.all(jnp.equal(f(idx).array, named1.array[0, :, :]))
+
+
 def test_index_array_slices():
     # fancier tests with array slices with named array args
     H = Axis("H", 10)

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

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