MAINT: deduplicate y != None in cov, corrcoef

Author: ev-br

torch_np/_detail/implementations.py

@@ -301,11 +301,7 @@ def vstack(tensors, *, dtype=None, casting="same_kind"):
 # #### 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
-
+def corrcoef(xy_tensor, *, dtype=None):
     is_half = dtype == torch.float16
     if is_half:
         # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"

torch_np/_wrapper.py

@@ -388,26 +388,40 @@ def diag(v, k=0):
 ###### misc/unordered
 
 
+def _xy_helper_corrcoef(x_tensor, y_tensor=None, rowvar=True):
+    """Prepate inputs for cov and corrcoef."""
+
+    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/function_base.py#L2636
+    if y_tensor is not None:
+        # make sure x and y are at least 2D
+        ndim_extra = 2 - x_tensor.ndim
+        if ndim_extra > 0:
+            x_tensor = x_tensor.view((1,) * ndim_extra + x_tensor.shape)
+        if not rowvar and x_tensor.shape[0] != 1:
+            x_tensor = x_tensor.mT
+        x_tensor = x_tensor.clone()
+
+        ndim_extra = 2 - y_tensor.ndim
+        if ndim_extra > 0:
+            y_tensor = y_tensor.view((1,) * ndim_extra + y_tensor.shape)
+        if not rowvar and y_tensor.shape[0] != 1:
+            y_tensor = y_tensor.mT
+        y_tensor = y_tensor.clone()
+
+        x_tensor = _impl.concatenate((x_tensor, y_tensor), axis=0)
+
+    return x_tensor
+
+
 @_decorators.dtype_to_torch
 def corrcoef(x, y=None, rowvar=True, bias=NoValue, ddof=NoValue, *, dtype=None):
     if bias is not None or ddof is not None:
         # deprecated in NumPy
         raise NotImplementedError
 
-    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/function_base.py#L2636
-    if y is not None:
-        x = array(x, ndmin=2)
-        if not rowvar and x.shape[0] != 1:
-            x = x.T
-
-        y = array(y, ndmin=2)
-        if not rowvar and y.shape[0] != 1:
-            y = y.T
-
-        x = concatenate((x, y), axis=0)
-
-    x_tensor = asarray(x).get()
-    result = _impl.corrcoef(x_tensor, rowvar, dtype=dtype)
+    x_tensor, y_tensor = _helpers.to_tensors_or_none(x, y)
+    tensor = _xy_helper_corrcoef(x_tensor, y_tensor, rowvar)
+    result = _impl.corrcoef(tensor, dtype=dtype)
     return asarray(result)
 
 
@@ -423,21 +437,12 @@ def cov(
     *,
     dtype=None,
 ):
-    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/function_base.py#L2636
-    if y is not None:
-        m = array(m, ndmin=2)
-        if not rowvar and m.shape[0] != 1:
-            m = m.T
-
-        y = array(y, ndmin=2)
-        if not rowvar and y.shape[0] != 1:
-            y = y.T
 
-        m = concatenate((m, y), axis=0)
-
-    m_tensor, fweights_tensor, aweights_tensor = _helpers.to_tensors_or_none(
-        m, fweights, aweights
+    m_tensor, y_tensor, fweights_tensor, aweights_tensor = _helpers.to_tensors_or_none(
+        m, y, fweights, aweights
     )
+    m_tensor = _xy_helper_corrcoef(m_tensor, y_tensor, rowvar)
+
     result = _impl.cov(
         m_tensor, bias, ddof, fweights_tensor, aweights_tensor, dtype=dtype
     )