Repository: jax-ml/jax
Pull request #27668 (2025-04-02 13:16:18)

`_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.
`_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.


georgevdd: Before:

```
>>> jax.make_jaxpr(lambda a: a[1:, ...])(jnp.ones([3]))
{ lambda ; a:f32[3]. let
    b:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] 1
    c:f32[2] = gather[
      dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
      fill_value=None
      indices_are_sorted=True
      mode=GatherScatterMode.PROMISE_IN_BOUNDS
      slice_sizes=(2,)
      unique_indices=True
    ] a b
  in (c,) }
```

After:

```
>>> jax.make_jaxpr(lambda a: a[1:, ...])(jnp.ones([3]))
{ lambda ; a:f32[3]. let
    b:f32[2] = slice[limit_indices=(3,) start_indices=(1,) strides=None] a
  in (b,) }
```

hawkinsp: > Before:
> 
> ```
> >>> jax.make_jaxpr(lambda a: a[1:, ...])(jnp.ones([3]))
> { lambda ; a:f32[3]. let
>     b:i32[1] = broadcast_in_dim[
>       broadcast_dimensions=()
>       shape=(1,)
>       sharding=None
>     ] 1
>     c:f32[2] = gather[
>       dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
>       fill_value=None
>       indices_are_sorted=True
>       mode=GatherScatterMode.PROMISE_IN_BOUNDS
>       slice_sizes=(2,)
>       unique_indices=True
>     ] a b
>   in (c,) }
> ```
> 
> After:
> 
> ```
> >>> jax.make_jaxpr(lambda a: a[1:, ...])(jnp.ones([3]))
> { lambda ; a:f32[3]. let
>     b:f32[2] = slice[limit_indices=(3,) start_indices=(1,) strides=None] a
>   in (b,) }
> ```


Seems reasonable, but if this is a property you may be depending on, is it worth a test?
georgevdd: You beat me to the proverbial punch while I was working out the best place to add the test :-)

georgevdd: I added test coverage. Is there anything else I need to do here, or should I just wait?

hawkinsp: > I added test coverage. Is there anything else I need to do here, or should I just wait?

Since you created this PR with our internal tooling, I think you might have to send it for review that way too.

Patch:
diff --git a/jax/_src/numpy/indexing.py b/jax/_src/numpy/indexing.py
index 863f0c775..05169dd54 100644
--- a/jax/_src/numpy/indexing.py
+++ b/jax/_src/numpy/indexing.py
@@ -526,8 +526,6 @@ def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) ->
 
   if not all(isinstance(i, int) for i in arr.shape):
     return None
-  if len(idx) > arr.ndim:
-    return None
   if any(i is None for i in idx):
     return None  # TODO(jakevdp): handle newaxis case
   # For symbolic dimensions fallback to gather
@@ -535,10 +533,13 @@ def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) ->
          for i in idx if isinstance(i, slice)
          for elt in (i.start, i.stop, i.step)):
     return None
-
   if any(i is Ellipsis for i in idx):
-    # Remove ellipses and add trailing `slice(None)`.
+    # Remove ellipses and pad with trailing `slice(None)` if necessary.
+    # Do this before checking against rank of `arr` so that `...` can
+    # count as no dimensions at all (e.g. `my_1d_array[:, ...]` succeeds)
     idx = _canonicalize_tuple_index(arr.ndim, idx=idx)
+  if len(idx) > arr.ndim:
+    return None
 
   simple_revs = {i for i, ind in enumerate(idx) if _is_simple_reverse_slice(ind)}
   int_indices = {i for i, (ind, size) in enumerate(zip(idx, arr.shape))


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index 63a725ad3..ca9ba9c88 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -926,12 +926,20 @@ class IndexingTest(jtu.JaxTestCase):
     self.assertEqual(jaxpr.jaxpr.eqns[-2].primitive, lax.slice_p)
     self.assertEqual(jaxpr.jaxpr.eqns[-1].primitive, lax.squeeze_p)
 
-    # Indexing with `Ellipsis` is not lowered to `gather`.
+    # Indexing with `Ellipsis` is not lowered to `gather` ...
     jaxpr = jax.make_jaxpr(lambda x: x[..., 0])(jnp.ones((3, 4, 5)))
     self.assertLen((jaxpr.jaxpr.eqns), 2)
     self.assertEqual(jaxpr.jaxpr.eqns[-2].primitive, lax.slice_p)
     self.assertEqual(jaxpr.jaxpr.eqns[-1].primitive, lax.squeeze_p)
 
+    # ... even when the ellipsis expands to no dimensions.
+    jaxpr = jax.make_jaxpr(lambda x: x[..., 0:1])(jnp.ones((3,)))
+    self.assertLen((jaxpr.jaxpr.eqns), 1)
+    self.assertEqual(jaxpr.jaxpr.eqns[-1].primitive, lax.slice_p)
+    jaxpr = jax.make_jaxpr(lambda x: x[0:1, ...])(jnp.ones((3,)))
+    self.assertLen((jaxpr.jaxpr.eqns), 1)
+    self.assertEqual(jaxpr.jaxpr.eqns[-1].primitive, lax.slice_p)
+
     # Simple reverses lower to lax.rev_p
     jaxpr = jax.make_jaxpr(lambda x: x[:, ::-1])(jnp.ones((3, 4)))
     self.assertEqual(len(jaxpr.jaxpr.eqns), 1)