MAINT: split cov, corrcoef, concat into ndarray/torch parts

Author: ev-br

torch_np/_decorators.py

@@ -15,7 +15,7 @@ def wrapped(*args, dtype=None, **kwds):
         torch_dtype = None
         if dtype is not None:
             dtype = _dtypes.dtype(dtype)
-            torch_dtype = dtype._scalar_type.torch_dtype
+            torch_dtype = dtype.torch_dtype
         return func(*args, dtype=torch_dtype, **kwds)
 
     return wrapped

torch_np/_detail/implementations.py

@@ -95,6 +95,82 @@ def diff(a_tensor, n=1, axis=-1, prepend_tensor=None, append_tensor=None):
     return result
 
 
+# #### concatenate and relatives
+
+def concatenate(tensors, axis=0, out=None, dtype=None, casting="same_kind"):
+    # np.concatenate ravels if axis=None
+    tensors, axis = _util.axis_none_ravel(*tensors, axis=axis)
+
+    # figure out the type of the inputs and outputs
+    if out is None and dtype is None:
+        out_dtype = None
+    else:
+        out_dtype = out.dtype.torch_dtype if dtype is None else dtype
+
+        # cast input arrays if necessary; do not broadcast them agains `out`
+        tensors = _util.cast_dont_broadcast(tensors, out_dtype, casting)
+
+    try:
+        result = torch.cat(tensors, axis)
+    except (IndexError, RuntimeError):
+        raise _util.AxisError
+
+    return result
+
+
+# #### cov & corrcoef
+
+def corrcoef(xy_tensor,  rowvar=True, *, dtype=None):
+    if rowvar is False:
+        # xy_tensor is at least 2D, so using .T is safe
+        xy_tensor = x_tensor.T
+
+    is_half = dtype == torch.float16
+    if is_half:
+        # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
+        dtype = torch.float32
+
+    if dtype is not None:
+        xy_tensor = xy_tensor.to(dtype)
+
+    result = torch.corrcoef(xy_tensor)
+
+    if is_half:
+        result = result.to(torch.float16)
+
+    return result
+
+
+def cov(
+    m_tensor,
+    bias=False,
+    ddof=None,
+    fweights_tensor=None,
+    aweights_tensor=None,
+    *,
+    dtype=None,
+):
+    if ddof is None:
+        ddof = 1 if bias == 0 else 0
+
+    is_half = dtype == torch.float16
+    if is_half:
+        # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
+        dtype = torch.float32
+
+    if dtype is not None:
+        m_tensor = m_tensor.to(dtype)
+
+    result = torch.cov(
+        m_tensor, correction=ddof, aweights=aweights_tensor, fweights=fweights_tensor
+    )
+
+    if is_half:
+        result = result.to(torch.float16)
+
+    return result
+
+
 def meshgrid(*xi_tensors, copy=True, sparse=False, indexing="xy"):
     # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/function_base.py#L4892-L5047
     ndim = len(xi_tensors)
@@ -118,3 +194,4 @@ def meshgrid(*xi_tensors, copy=True, sparse=False, indexing="xy"):
         output = [x.clone() for x in output]
 
     return output
+

torch_np/_wrapper.py

@@ -434,23 +434,7 @@ def corrcoef(x, y=None, rowvar=True, bias=NoValue, ddof=NoValue, *, dtype=None):
         x = concatenate((x, y), axis=0)
 
     x_tensor = asarray(x).get()
-
-    if rowvar is False:
-        x_tensor = x_tensor.T
-
-    is_half = dtype == torch.float16
-    if is_half:
-        # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
-        dtype = torch.float32
-
-    if dtype is not None:
-        x_tensor = x_tensor.to(dtype)
-
-    result = torch.corrcoef(x_tensor)
-
-    if is_half:
-        result = result.to(torch.float16)
-
+    result = _impl.corrcoef(x_tensor, rowvar, dtype=dtype)
     return asarray(result)
 
 
@@ -478,45 +462,25 @@ def cov(
 
         m = concatenate((m, y), axis=0)
 
-    if ddof is None:
-        if bias == 0:
-            ddof = 1
-        else:
-            ddof = 0
+#    if ddof is None:
+#        if bias == 0:
+#            ddof = 1
+#        else:
+#            ddof = 0
 
     m_tensor, fweights_tensor, aweights_tensor = _helpers.to_tensors_or_none(
         m, fweights, aweights
     )
-
-    # work with tensors from now on
-    is_half = dtype == torch.float16
-    if is_half:
-        # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
-        dtype = torch.float32
-
-    if dtype is not None:
-        m_tensor = m_tensor.to(dtype)
-
-    result = torch.cov(
-        m_tensor, correction=ddof, aweights=aweights_tensor, fweights=fweights_tensor
-    )
-
-    if is_half:
-        result = result.to(torch.float16)
-
+    result = _impl.cov(m_tensor, bias, ddof, fweights_tensor, aweights_tensor, dtype=dtype)
     return asarray(result)
 
 
+@_decorators.dtype_to_torch
 def concatenate(ar_tuple, axis=0, out=None, dtype=None, casting="same_kind"):
     if ar_tuple == ():
         # XXX: RuntimeError in torch, ValueError in numpy
         raise ValueError("need at least one array to concatenate")
 
-    tensors = _helpers.to_tensors(*ar_tuple)
-
-    # np.concatenate ravels if axis=None
-    tensors, axis = _util.axis_none_ravel(*tensors, axis=axis)
-
     if out is not None:
         if not isinstance(out, ndarray):
             raise ValueError("'out' must be an array")
@@ -527,22 +491,8 @@ def concatenate(ar_tuple, axis=0, out=None, dtype=None, casting="same_kind"):
                 "concatenate() only takes `out` or `dtype` as an "
                 "argument, but both were provided."
             )
-
-    # figure out the type of the inputs and outputs
-    if out is None and dtype is None:
-        out_dtype = None
-    else:
-        out_dtype = out.dtype if dtype is None else _dtypes.dtype(dtype)
-        out_dtype = out_dtype.type.torch_dtype
-
-        # cast input arrays if necessary; do not broadcast them agains `out`
-        tensors = _util.cast_dont_broadcast(tensors, out_dtype, casting)
-
-    try:
-        result = torch.cat(tensors, axis)
-    except (IndexError, RuntimeError):
-        raise _util.AxisError
-
+    tensors = _helpers.to_tensors(*ar_tuple)
+    result = _impl.concatenate(tensors, axis, out, dtype, casting)
     return _helpers.result_or_out(result, out)