take() implementation

Author: honno

torch_np/_funcs.py

@@ -929,6 +929,21 @@ def asfarray():
 
 # ### put/take_along_axis ###
 
+@normalizer
+def take(a: ArrayLike, indices: ArrayLike, axis=None, out : Optional[NDArray]=None, mode="raise"):
+    if out is not None:
+        raise NotImplementedError(f"{out=}")
+    if mode != "raise":
+        raise NotImplementedError(f"{mode=}")
+
+    (a,), axis = _util.axis_none_ravel(a, axis=axis)
+    axis = _util.normalize_axis_index(axis, a.ndim)
+    idx = (slice(None),) * axis + (indices, ...)
+    result = a[idx]
+    return result
+
+
+
 
 @normalizer
 def take_along_axis(arr: ArrayLike, indices: ArrayLike, axis):

torch_np/_ndarray.py

@@ -407,8 +407,7 @@ def __setitem__(self, index, value):
         value = _helpers.ndarrays_to_tensors(value)
         return self.tensor.__setitem__(index, value)
 
-    def take(*a, **kw):
-        raise NotImplementedError()
+    take = _funcs.take
 
 
 # This is the ideally the only place which talks to ndarray directly.