BUG: Fixed an issue where `nd_grid` could cast `float128` to `float64`

1 files changed, 20 insertions, 17 deletions

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index 2a4402c89..7ef312dbb 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -148,24 +148,25 @@ class nd_grid:
     def __getitem__(self, key):
         try:
             size = []
-            typ = int
-            for kk in key:
-                step = kk.step
-                start = kk.start
+            # Mimic the behavior of `np.arange` and use a data type
+            # which is at least as large as `np.int_`
+            num_list = [0]
+            for k in range(len(key)):
+                step = key[k].step
+                start = key[k].start
+                stop = key[k].stop
                 if start is None:
                     start = 0
                 if step is None:
                     step = 1
                 if isinstance(step, (_nx.complexfloating, complex)):
-                    size.append(int(abs(step)))
-                    typ = float
+                    step = abs(step)
+                    size.append(int(step))
                 else:
                     size.append(
-                        int(math.ceil((kk.stop - start) / (step * 1.0))))
-                if (isinstance(step, (_nx.floating, float)) or
-                        isinstance(start, (_nx.floating, float)) or
-                        isinstance(kk.stop, (_nx.floating, float))):
-                    typ = float
+                        int(math.ceil((stop - start) / (step*1.0))))
+                num_list += [start, stop, step]
+            typ = _nx.result_type(*num_list)
             if self.sparse:
                 nn = [_nx.arange(_x, dtype=_t)
                       for _x, _t in zip(size, (typ,)*len(size))]
@@ -197,11 +198,13 @@ class nd_grid:
             if start is None:
                 start = 0
             if isinstance(step, (_nx.complexfloating, complex)):
-                step = abs(step)
-                length = int(step)
+                # Prevent the (potential) creation of integer arrays
+                step_float = abs(step)
+                step = length = int(step_float)
                 if step != 1:
                     step = (key.stop-start)/float(step-1)
-                return _nx.arange(0, length, 1, float)*step + start
+                typ = _nx.result_type(start, stop, step_float)
+                return _nx.arange(0, length, 1, dtype=typ)*step + start
             else:
                 return _nx.arange(start, stop, step)
 
@@ -621,7 +624,7 @@ class ndindex:
     Parameters
     ----------
     shape : ints, or a single tuple of ints
-        The size of each dimension of the array can be passed as 
+        The size of each dimension of the array can be passed as
         individual parameters or as the elements of a tuple.
 
     See Also
@@ -631,7 +634,7 @@ class ndindex:
     Examples
     --------
     Dimensions as individual arguments
-    
+
     >>> for index in np.ndindex(3, 2, 1):
     ...     print(index)
     (0, 0, 0)
@@ -642,7 +645,7 @@ class ndindex:
     (2, 1, 0)
 
     Same dimensions - but in a tuple ``(3, 2, 1)``
-    
+
     >>> for index in np.ndindex((3, 2, 1)):
     ...     print(index)
     (0, 0, 0)