MAINT: address review comments

Author: ev-br

torch_np/_detail/implementations.py

@@ -31,7 +31,7 @@ def tensor_equiv(a1_t, a2_t):
     return tensor_equal(a1_t, a2_t)
 
 
-def tensor_isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
+def isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
     dtype = _dtypes_impl.result_type_impl((a.dtype, b.dtype))
     a = a.to(dtype)
     b = b.to(dtype)
@@ -42,7 +42,7 @@ def tensor_isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
 # ### is arg real or complex valued ###
 
 
-def tensor_iscomplex(x):
+def iscomplex(x):
     if torch.is_complex(x):
         return torch.as_tensor(x).imag != 0
     result = torch.zeros_like(x, dtype=torch.bool)
@@ -51,7 +51,7 @@ def tensor_iscomplex(x):
     return result
 
 
-def tensor_isreal(x):
+def isreal(x):
     if torch.is_complex(x):
         return torch.as_tensor(x).imag == 0
     result = torch.ones_like(x, dtype=torch.bool)
@@ -60,7 +60,8 @@ def tensor_isreal(x):
     return result
 
 
-def tensor_real_if_close(x, tol=100):
+def real_if_close(x, tol=100):
+    # XXX: copies vs views; numpy seems to return a copy?
     if not torch.is_complex(x):
         return x
     mask = torch.abs(x.imag) < tol * torch.finfo(x.dtype).eps
@@ -73,20 +74,20 @@ def tensor_real_if_close(x, tol=100):
 # ### math functions ###
 
 
-def tensor_angle(z, deg=False):
+def angle(z, deg=False):
     result = torch.angle(z)
     if deg:
-        result *= 180 / torch.pi
+        result = result * 180 / torch.pi
     return result
 
 
 # ### sorting ###
 
 
-def tensor_argsort(tensor, axis=-1, kind=None, order=None):
+def argsort(tensor, axis=-1, kind=None, order=None):
     if order is not None:
         raise NotImplementedError
-    stable = True if kind == "stable" else False
+    stable = kind == "stable"
     if axis is None:
         axis = -1
     return torch.argsort(tensor, stable=stable, dim=axis, descending=False)
@@ -387,11 +388,11 @@ def bincount(x_tensor, /, weights_tensor=None, minlength=0):
 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))
+    base = torch.pow(stop / start, 1.0 / (num - 1))
+    logbase = torch.log(base)
     result = torch.logspace(
-        torch.log(tstart) / torch.log(base),
-        torch.log(tstop) / torch.log(base),
+        torch.log(start) / logbase,
+        torch.log(stop) / logbase,
         num,
         base=base,
     )

torch_np/_wrapper.py

@@ -261,6 +261,7 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis
 def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     if axis != 0 or not endpoint:
         raise NotImplementedError
+    start, stop = _helpers.to_tensors(start, stop)
     result = _impl.geomspace(start, stop, num, endpoint, dtype, axis)
     return asarray(result)
 
@@ -954,7 +955,7 @@ def diff(a, n=1, axis=-1, prepend=NoValue, append=NoValue):
 
 @asarray_replacer()
 def argsort(a, axis=-1, kind=None, order=None):
-    result = _impl.tensor_argsort(a, axis, kind, order)
+    result = _impl.argsort(a, axis, kind, order)
     return result
 
 
@@ -963,7 +964,7 @@ def argsort(a, axis=-1, kind=None, order=None):
 
 @asarray_replacer()
 def angle(z, deg=False):
-    result = _impl.tensor_angle(z, deg)
+    result = _impl.angle(z, deg)
     return result
 
 
@@ -984,19 +985,19 @@ def imag(a):
 
 @asarray_replacer()
 def real_if_close(a, tol=100):
-    result = _impl.tensor_real_if_close(a, tol=tol)
+    result = _impl.real_if_close(a, tol=tol)
     return result
 
 
 @asarray_replacer()
 def iscomplex(x):
-    result = _impl.tensor_iscomplex(x)
+    result = _impl.iscomplex(x)
     return result  # XXX: missing .item on a zero-dim value; a case for array_or_scalar(value) ?
 
 
 @asarray_replacer()
 def isreal(x):
-    result = _impl.tensor_isreal(x)
+    result = _impl.isreal(x)
     return result
 
 
@@ -1036,13 +1037,13 @@ def isscalar(a):
 
 def isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
     a_t, b_t = _helpers.to_tensors(a, b)
-    result = _impl.tensor_isclose(a_t, b_t, rtol, atol, equal_nan=equal_nan)
+    result = _impl.isclose(a_t, b_t, rtol, atol, equal_nan=equal_nan)
     return asarray(result)
 
 
 def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
     a_t, b_t = _helpers.to_tensors(a, b)
-    result = _impl.tensor_isclose(a_t, b_t, rtol, atol, equal_nan=equal_nan)
+    result = _impl.isclose(a_t, b_t, rtol, atol, equal_nan=equal_nan)
     return result.all()