Repository: google/jax

Issue:
Please make jnp.delete jittable.
Current workaround:
```
@jax.jit
def apply_except_one(x, i):
  idx = jnp.arange(x.shape[0])
  idx = jnp.where(idx >= i, idx + 1, idx)[:-1]
  return _complex_jittable_function(x[idx])
```

Desired state:
```
@jax.jit
def apply_except_one(x, i):
  return _complex_jittable_function(jnp.delete(x, i))
```


Discussion:
```
>>> jax.jit(jnp.delete, static_argnums=1)(jnp.arange(5), 2)
DeviceArray([0, 1, 3, 4], dtype=int32)
```

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError and https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit
Sorry, there's chat context here! We don't want to use `static_argnums`, and indeed we don't really have to. (See the JAXers thread with Blake.)
Just saw that thread :)

Assigning to you instead :)
Assign Peter to everything, like we talked about! ;)



> Sorry, there's chat context here! We don't want to use `static_argnums`, and indeed we don't really have to. (See the JAXers thread with Blake.)

@mattjj Hi matthew, I think I miss the chat or thread ? so how exactly make jnp.delete jittable without static_argnums? :)

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2883,8 +2883,18 @@ def append(arr, values, axis: Optional[int] = None):
     return concatenate([arr, values], axis=axis)
 
 
-@util._wraps(np.delete)
-def delete(arr, obj, axis=None):
+@util._wraps(np.delete,
+  lax_description=_dedent("""
+    delete() usually requires the index specification to be static. If the index
+    is an integer array that is guaranteed to contain unique entries, you may
+    specify ``assume_unique_indices=True`` to perform the operation in a
+    manner that does not require static indices."""),
+  extra_params=_dedent("""
+    assume_unique_indices : int, optional (default=False)
+        In case of array-like integer (not boolean) indices, assume the indices are unique,
+        and perform the deletion in a way that is compatible with JIT and other JAX
+        transformations."""))
+def delete(arr, obj, axis=None, *, assume_unique_indices=False):
   util.check_arraylike("delete", arr)
   if axis is None:
     arr = ravel(arr)
@@ -2910,6 +2920,18 @@ def delete(arr, obj, axis=None):
   # Case 3: obj is an array
   # NB: pass both arrays to check for appropriate error message.
   util.check_arraylike("delete", arr, obj)
+
+  # Case 3a: unique integer indices; delete in a JIT-compatible way
+  if issubdtype(_dtype(obj), integer) and assume_unique_indices:
+    obj = asarray(obj).ravel()
+    obj = clip(where(obj < 0, obj + arr.shape[axis], obj), 0, arr.shape[axis])
+    obj = sort(obj)
+    obj -= arange(len(obj))
+    i = arange(arr.shape[axis] - obj.size)
+    i += (i[None, :] >= obj[:, None]).sum(0)
+    return arr[(slice(None),) * axis + (i,)]
+
+  # Case 3b: non-unique indices: must be static.
   obj = core.concrete_or_error(np.asarray, obj, "'obj' array argument of jnp.delete()")
 
   if issubdtype(obj.dtype, integer):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1526,6 +1526,29 @@ def testDeleteIndexArray(self, shape, dtype, axis, idx_shape):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @jtu.sample_product(
+    [dict(shape=shape, axis=axis)
+      for shape in nonempty_nonscalar_array_shapes
+      for axis in [None] + list(range(-len(shape), len(shape)))
+    ],
+    dtype=all_dtypes,
+    idx_shape=all_shapes,
+  )
+  def testDeleteUniqueIndices(self, shape, dtype, axis, idx_shape):
+    rng = jtu.rand_default(self.rng())
+    max_idx = np.zeros(shape).size if axis is None else np.zeros(shape).shape[axis]
+    idx_size = np.zeros(idx_shape).size
+    if idx_size > max_idx:
+      self.skipTest("Too many indices to be unique")
+    def args_maker():
+      x = rng(shape, dtype)
+      idx = self.rng().choice(max_idx, idx_shape, replace=False)
+      return x, idx
+    np_fun = partial(np.delete, axis=axis)
+    jnp_fun = partial(jnp.delete, axis=axis, assume_unique_indices=True)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @jtu.sample_product(
     [dict(shape=shape, axis=axis)
       for shape in nonempty_nonscalar_array_shapes