MAINT: IndexExpression into _funcs.py from _detail/_index_tricks.py

Author: ev-br

torch_np/__init__.py

@@ -1,6 +1,5 @@
 from . import random
 from ._binary_ufuncs import *
-from ._detail._index_tricks import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
 from ._funcs import *

torch_np/_detail/_index_tricks.py

@@ -1871,3 +1871,30 @@ def isscalar(a: ArrayLike):
     if a is False:
         return a
     return a.numel() == 1
+
+
+"""
+Vendored objects from numpy.lib.index_tricks
+"""
+
+
+class IndexExpression:
+    """
+    Written by Konrad Hinsen <[email protected]>
+    last revision: 1999-7-23
+
+    Cosmetic changes by T. Oliphant 2001
+    """
+
+    def __init__(self, maketuple):
+        self.maketuple = maketuple
+
+    def __getitem__(self, item):
+        if self.maketuple and not isinstance(item, tuple):
+            return (item,)
+        else:
+            return item
+
+
+index_exp = IndexExpression(maketuple=True)
+s_ = IndexExpression(maketuple=False)

torch_np/_funcs.py

@@ -13,7 +13,7 @@
     )
 
 from torch_np import diag_indices, diag_indices_from, fill_diagonal
-from torch_np._detail._index_tricks import index_exp, s_
+from torch_np import index_exp, s_
 
 
 @pytest.mark.xfail(reason='unravel_index not implemented')