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

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 37d27140e..6e4e44b80 100644
--- 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):
['jax/_src/numpy/lax_numpy.py']