Merge pull request #7907 from jakevdp:take-arraylike

PiperOrigin-RevId: 396584251
google · Sep 14, 2021 · 788ecc5 · 788ecc5
2 parents c365d7f + 066a3bd
commit 788ecc5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -5079,6 +5079,7 @@ def take(a, indices, axis: Optional[int] = None, out=None, mode=None):
 def _take(a, indices, axis: Optional[int] = None, out=None, mode=None):
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.take is not supported.")
+  _check_arraylike("take", a)
   a = asarray(a)
   indices = asarray(indices)