Repository: google/jax

Issue:
Jax numpy array_split is slow compared to numpy
I was writing Jax code, substituting all of the places where I would normally use NumPy with jax.numpy
In using jax.numpy.array_split with a large input, I found that the program was incredibly slow (for jax+numpy or jax then the program didn't finish in 5 minutes for me) 

While running there was no issue however if I ctrl+c to terminate then the follow error was thrown
`F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:446] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: ptxas exited with non-zero error code 2, output: '  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Aborted (core dumped)`

It seems that Jax is trying to jit/compile the array input which is slow due to the size, I just wanted to check if this was expected or not
If this is expected, could some documentation be added to note this problem on the relative functions 

Example code
```
from tqdm import tqdm
import numpy as onp
import jax.numpy as jnp

print('numpy')
a = [batch for batch in tqdm(onp.array_split(onp.arange(10_000), 1_000))]
b = [batch for batch in tqdm(onp.array_split(onp.arange(100_000), 10_000))]

print('jax + numpy')
c = [batch for batch in tqdm(jnp.array_split(onp.arange(100_000), 10_000))]
d = [batch for batch in tqdm(jnp.array_split(onp.arange(100_000), 10_000))]

print('jax')
e = [batch for batch in tqdm(jnp.array_split(jnp.arange(100_000), 10_000))]
f = [batch for batch in tqdm(jnp.array_split(jnp.arange(100_000), 10_000))]
```

Thanks for any help



Discussion:
Thanks for the report – I think the issue is that `jnp.array_split` fails to push the numpy input to device before processing it, meaning that the splitting is effectively done on client side and then each output is pushed to device individually, incurring that overhead thousands of times instead of one time.

I'll work on a fix.
@jakevdp It's also possible that that code would be faster if it used `dynamic_slice` rather than `slice`, especially when not under a `jit`. Otherwise we need to compile one "slice" operator for each index. In a `jit` it will make no difference.
Good point - my first PR gives an order of magnitude improvement by calling device_put, so I'll get that in first and then take a look at `dynamic_slice`.

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
@@ -2260,6 +2260,8 @@ def broadcast_to(arr, shape):
 
 
 def _split(op, ary, indices_or_sections, axis=0):
+  _check_arraylike(op, ary)
+  ary = asarray(ary)
   axis = core.concrete_or_error(int, axis, f"in jax.numpy.{op} argument `axis`")
   size = ary.shape[axis]
   if isinstance(indices_or_sections, (tuple, list)):
@@ -2298,15 +2300,15 @@ def _split(op, ary, indices_or_sections, axis=0):
 def split(ary, indices_or_sections, axis: int = 0):
   return _split("split", ary, indices_or_sections, axis=axis)
 
-def _split_on_axis(np_fun, axis):
-  @_wraps(np_fun, update_doc=False)
+def _split_on_axis(op, axis):
+  @_wraps(getattr(np, op), update_doc=False)
   def f(ary, indices_or_sections):
-    return split(ary, indices_or_sections, axis=axis)
+    return _split(op, ary, indices_or_sections, axis=axis)
   return f
 
-vsplit = _split_on_axis(np.vsplit, axis=0)
-hsplit = _split_on_axis(np.hsplit, axis=1)
-dsplit = _split_on_axis(np.dsplit, axis=2)
+vsplit = _split_on_axis("vsplit", axis=0)
+hsplit = _split_on_axis("hsplit", axis=1)
+dsplit = _split_on_axis("dsplit", axis=2)
 
 @_wraps(np.array_split)
 def array_split(ary, indices_or_sections, axis: int = 0):


Test patch: