Vendor s_ and index_exp

Author: honno

torch_np/__init__.py

@@ -1,6 +1,7 @@
 from ._wrapper import *  # isort: skip  # XXX: currently this prevents circular imports
 from . import random
 from ._binary_ufuncs import *
+from ._detail._index_tricks import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
 from ._getlimits import errstate, finfo, iinfo
@@ -16,8 +17,5 @@
 nan = float("nan")
 from math import pi  # isort: skip
 
-from numpy import s_  # isort: skip
-from numpy import index_exp  # isort: skip
-
 False_ = asarray(False, bool_)
 True_ = asarray(True, bool_)

torch_np/_detail/_index_tricks.py

@@ -0,0 +1,26 @@
+"""
+Vendored objects from numpy.lib.index_tricks
+"""
+__all__ = ["index_exp", "s_"]
+
+
+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)