SQUASH: premiliminary split arange/geomspace

Author: ev-br

torch_np/_detail/implementations.py

@@ -138,7 +138,8 @@ def split_helper(tensor, indices_or_sections, axis, strict=False):
     if isinstance(indices_or_sections, int):
         return split_helper_int(tensor, indices_or_sections, axis, strict)
     elif isinstance(indices_or_sections, (list, tuple)):
-        return split_helper_list(tensor, list(indices_or_sections), axis, strict)
+        # NB: drop split=..., it only applies to split_helper_int
+        return split_helper_list(tensor, list(indices_or_sections), axis)
     else:
         raise TypeError("split_helper: ", type(indices_or_sections))
 
@@ -170,7 +171,7 @@ def split_helper_int(tensor, indices_or_sections, axis, strict=False):
     return result
 
 
-def split_helper_list(tensor, indices_or_sections, axis, strict=False):
+def split_helper_list(tensor, indices_or_sections, axis):
     if not isinstance(indices_or_sections, list):
         raise NotImplementedError("split: indices_or_sections: list")
     # numpy expectes indices, while torch expects lengths of sections
@@ -381,3 +382,44 @@ def bincount(x_tensor, /, weights_tensor=None, minlength=0):
 
     result = torch.bincount(x_tensor, weights_tensor, minlength)
     return result
+
+
+# ### linspace, geomspace, logspace and arange ###
+
+def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
+    if axis != 0 or not endpoint:
+        raise NotImplementedError
+    tstart, tstop = torch.as_tensor([start, stop])
+    base = torch.pow(tstop / tstart, 1.0 / (num - 1))
+    result = torch.logspace(
+        torch.log(tstart) / torch.log(base),
+        torch.log(tstop) / torch.log(base),
+        num,
+        base=base,
+    )
+    return result
+
+
+
+def arange(start=None, stop=None, step=1, dtype=None):
+    if step == 0:
+        raise ZeroDivisionError
+    if stop is None and start is None:
+        raise TypeError
+    if stop is None:
+        # XXX: this breaks if start is passed as a kwarg:
+        # arange(start=4) should raise (no stop) but doesn't
+        start, stop = 0, start
+    if start is None:
+        start = 0
+
+    if dtype is None:
+        dt_list = [_util._coerce_to_tensor(x).dtype for x in (start, stop, step)]
+        dtype = _dtypes_impl.default_int_dtype
+        dt_list.append(dtype)
+        dtype = _dtypes_impl.result_type_impl(dt_list)
+
+    try:
+        return torch.arange(start, stop, step, dtype=dtype)
+    except RuntimeError:
+        raise ValueError("Maximum allowed size exceeded")

torch_np/_wrapper.py

@@ -257,49 +257,27 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis
     return asarray(torch.linspace(start, stop, num, dtype=dtype))
 
 
+@_decorators.dtype_to_torch
 def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     if axis != 0 or not endpoint:
         raise NotImplementedError
-    tstart, tstop = torch.as_tensor([start, stop])
-    base = torch.pow(tstop / tstart, 1.0 / (num - 1))
-    result = torch.logspace(
-        torch.log(tstart) / torch.log(base),
-        torch.log(tstop) / torch.log(base),
-        num,
-        base=base,
-    )
+    result = _impl.geomspace(start, stop, num, endpoint, dtype, axis)
     return asarray(result)
 
 
+@_decorators.dtype_to_torch
 def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
     if axis != 0 or not endpoint:
         raise NotImplementedError
     return asarray(torch.logspace(start, stop, num, base=base, dtype=dtype))
 
 
+@_decorators.dtype_to_torch
 def arange(start=None, stop=None, step=1, dtype=None, *, like=None):
     _util.subok_not_ok(like)
-    if step == 0:
-        raise ZeroDivisionError
-    if stop is None and start is None:
-        raise TypeError
-    if stop is None:
-        # XXX: this breaks if start is passed as a kwarg:
-        # arange(start=4) should raise (no stop) but doesn't
-        start, stop = 0, start
-    if start is None:
-        start = 0
-
-    if dtype is None:
-        dtype = _dtypes.default_int_type()
-        dtype = result_type(start, stop, step, dtype)
-    torch_dtype = _dtypes.torch_dtype_from(dtype)
     start, stop, step = _helpers.ndarrays_to_tensors(start, stop, step)
-
-    try:
-        return asarray(torch.arange(start, stop, step, dtype=torch_dtype))
-    except RuntimeError:
-        raise ValueError("Maximum allowed size exceeded")
+    result = _impl.arange(start, stop, step, dtype=dtype)
+    return asarray(result)
 
 
 @_decorators.dtype_to_torch