Repository: google/jax

Issue:
Improve jnp.split API
Sorry for the very uninformative title for this issue, couldn't think of a better one.

While reading the code for `split` in `jax/_src/numpy/lax_numpy.py`, I noticed the following block of code for the function `lax_numpy._split`:

https://github.com/google/jax/blob/2b716f292d814fdfadca43e108b5fedc122c0d2c/jax/_src/numpy/lax_numpy.py#L1178-L1193

Accounting to the type signature of `indices_or_sections`

https://github.com/google/jax/blob/2b716f292d814fdfadca43e108b5fedc122c0d2c/jax/_src/numpy/lax_numpy.py#L1172-L1173

We should be able to remove the first branch (i.e. `(tuple, list)`), in the spirit of #7737. Was about to send a pull request. But I also noticed that `indices_or_sections` is always concrete, and `Array`s are not hashable. Keeping the current type also prevents us from marking `indices_or_sections` as static in `jit`. One might argue removing the second branch(i.e. `(np.ndarray, ndarray)`) and make `indices_or_sections: int | Sequence[int]` made more sense, but that's a deviation from the original numpy api.

Not sure what's the best action to take here.


Discussion:
I suspect it's the type annotations that are incorrect, not the implementation. The latter has much more test coverage than the former.
I would hesitate to apply #7737 here, as the argument must be static and so there's no performance footgun in passing a list instead of an array, and erroring on list/tuple input would break a lot of downstream code.
> I suspect it's the type annotations that are incorrect, not the implementation.

Totally agree. Can we error on array input and mark `split` as `static`?
Erroring on array input would be an API change that would require a deprecation cycle and also potentially a large number of fixes to downstream packages, and I can't think of any advantage that such a depracation would bring about (aside from one less conditional block in the impelmentation), so I lean toward no.
Should we at least relax the type annotation then? Feel free to close this issue.

> I can't think of any advantage that such a deprecation would bring about.

Even with the relax typing, use of `split` under `jit` still requires special treatment (other `jnp` functions will at least work under `jit` with correct `static_argnums`). Just pointing out the awkwardness, won't argue too strongly about it.

```python
a: Array
i: Array

def f(a, i):
  ...
  r = split(a, i)
  ...

f(a, i)                                   # type checks, works, let me try to jit it
jit(f)(a, i)                              # doesn't work, one could argue `i` should be marked as static
jit(f, static_argnums=(1,))(a, i)         # still doesn't work
jit(f, static_argnums=(1,))(a, tuple(i))  # finally
```
Yeah, understood that this does not work with JIT, and I think that's handled about as well as it can be in the code (explicit concreteness checks with associated error messages). We could probably adjust the type annotation to make clear that sequences are supported - what do you think?
> I think that's handled about as well as it can be in the code

That's fair, and agreed. Let's just change the annotation.

As a general question, should one expect `jit` (with correct `static_argnums` marking, for naked `jit` are known to change strictness of arguments) to change the apparent ~strictness~ staticness of some arguments of `f` for `f` in `jax.numpy`?
Can you elaborate on what you mean by "apparent strictness"?
Ah, there was a typo, I meant staticness. Or, maybe "appeared un-staticness".
<details>
Staticness is enforced at runtime in two ways:
- E1: Trying to concretise, example: in native control-flow
- E2: Disallow certain types, example: `jit(jnp.sum, static_argnums=(1,))(a, axis=jnp.int_(0))` errors out


Inside `jit`, (E1) catch error related to other aspect of staticness, like the one in `jit(jnp.nonzero)(a)`, i.e. some extra information needed, but the argument itself remains 'dynamic'(no type conversion).

Outside of `jit`, one usually only have (E1), inside of `jit` (with correct staticness  marking), both. Knowing (E2), supposedly, I could deduce staticness (or rather un-staticness), from type annotation. But it's 'appeared' for sometimes it is really only enforced by (1) at runtime:

```python
jnp.sum(a, axis=jnp.int_(0))  # runs, but type error
```

We could debug problem regarding staticness of a function `f` traversing following graph breadth-first, and would ideally back out before reaching state `d`:

```mermaid
flowchart TB
    A["a. Use knowledge about it's functionality"]
    B["b. Plug in some arrays and check for static type error (E2)"]
    C["c. Wrap in jit and check for runtime error (E1)"]
    D["d. Wrap the *outer* function with jit(static_argnums) and brute force"]
    A-->B
    A-->C
    B-->D
    subgraph jitted
    C-->D
    end
    B-->C
```

For most functions in `jax.numpy`, following `a → b` is enough, and only occasionally, `a → c`. Noticeably, even if `b` checks out and `c` doesn't (`a → b → c`, think `jnp.nonzero`), the fix in the end doesn't involve type conversion which is inherently contradict to the knowledge we aquired at state `b`.

The relaxed annotation for `jnp.split` forces us to reach `d`, and the staticness one *can* infer from the type is different from the real one acquires at the end. One might as well conclude it's `jit` that changes the staticness of argument `indices_or_sections`.

Of course, trade-offs, ease-of-use, python-typing-weak, etc. So I don't stand my ground 100%. Just really want to try out the flowchart functionality. :)
<details/>
> As a general question, should one expect `jit` (with correct `static_argnums` marking, for naked `jit` are known to change strictness of arguments) to change the apparent ~strictness~ staticness of some arguments of `f` for `f` in `jax.numpy`?

I believe `jit` will always cause a function's arguments to be non-static unless they're explicitly marked static using `static_argnames` or `static_argnums`. We might imagine a world where JIT tries to infer whether arguments should be static or traced, but JAX's current JIT implementation does not work that way.
I ran into this issue. I want to pass a `list` of plain Python `int`s to a `jnp.split` operation inside a `jit`. (The list is constructed by a complex series of function calls.) If I pass the list directly, I get a mypy error:
```
error: Argument 2 to "split" has incompatible type "List[int]"; expected "Union[int, Union[Array, ndarray[Any, Any], bool_, number[Any], bool, int, float, complex]]"  [arg-type]
```
If I apply `jnp.array` to the list, I get a runtime error:
```
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=6/0)>
in jax.numpy.split argument 1
```
What's the recommended solution to this? Is there any plan to relax the type annotation of `jnp.split` to allow `Sequence`s? For now, I'm using `typing.cast(jax.Array, indices)` as a workaround to satisfy mypy.
Try passing `np.array(indices)` instead: this will type-check as `ArrayLike` but remain static.
But I think the type annotation is wrong on `jnp.split`, it should be `indices_or_sections: Union[int, ArrayLike, Sequence[int]]` I think. The implementation explicitly handles `tuple` or `list` of ints.
@jakevdp I think so, too. I'm looking at https://github.com/google/jax/blob/3c1f3abba2c2e9b2506dc6bd19a132f60e2c320b/jax/_src/numpy/lax_numpy.py#L1163

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
@@ -1154,7 +1154,8 @@ def broadcast_to(array: ArrayLike, shape: Shape) -> Array:
   return util._broadcast_to(array, shape)
 
 
-def _split(op: str, ary: ArrayLike, indices_or_sections: Union[int, ArrayLike],
+def _split(op: str, ary: ArrayLike,
+           indices_or_sections: Union[int, Sequence[int], ArrayLike],
            axis: int = 0) -> List[Array]:
   util.check_arraylike(op, ary)
   ary = asarray(ary)
@@ -1193,12 +1194,13 @@ def _split(op: str, ary: ArrayLike, indices_or_sections: Union[int, ArrayLike],
           for start, end in zip(split_indices[:-1], split_indices[1:])]
 
 @util._wraps(np.split, lax_description=_ARRAY_VIEW_DOC)
-def split(ary: ArrayLike, indices_or_sections: Union[int, ArrayLike], axis: int = 0) -> List[Array]:
+def split(ary: ArrayLike, indices_or_sections: Union[int, Sequence[int], ArrayLike],
+          axis: int = 0) -> List[Array]:
   return _split("split", ary, indices_or_sections, axis=axis)
 
 def _split_on_axis(op: str, axis: int) -> Callable[[ArrayLike, Union[int, ArrayLike]], List[Array]]:
   @util._wraps(getattr(np, op), update_doc=False)
-  def f(ary: ArrayLike, indices_or_sections: Union[int, ArrayLike]) -> List[Array]:
+  def f(ary: ArrayLike, indices_or_sections: Union[int, Sequence[int], ArrayLike]) -> List[Array]:
     # for 1-D array, hsplit becomes vsplit
     nonlocal axis
     util.check_arraylike(op, ary)
@@ -1213,7 +1215,8 @@ def f(ary: ArrayLike, indices_or_sections: Union[int, ArrayLike]) -> List[Array]
 dsplit = _split_on_axis("dsplit", axis=2)
 
 @util._wraps(np.array_split)
-def array_split(ary: ArrayLike, indices_or_sections: Union[int, ArrayLike], axis: int = 0) -> List[Array]:
+def array_split(ary: ArrayLike, indices_or_sections: Union[int, Sequence[int], ArrayLike],
+                axis: int = 0) -> List[Array]:
   return _split("array_split", ary, indices_or_sections, axis=axis)
 
 @util._wraps(np.clip, skip_params=['out'])


Test patch: