[Array API] Add linalg.vecdot

This PR adds the function `linalg.vecdot` specified by the [Array
API](https://data-apis.org/array-api/latest/API_specification/linear_algebra_functions.html#function-vecdot)

For the complex case, it chooses to implement \sum x_i y_i. See the
discussion in data-apis/array-api#356

[ghstack-poisoned]

Author: lezcano

aten/src/ATen/core/interned_strings.h

@@ -218,6 +218,7 @@ namespace c10 {
   _(aten, mH)                        \
   _(aten, linalg_matrix_power)       \
   _(aten, chain_matmul)              \
+  _(aten, linalg_vecdot)             \
   _(aten, linalg_multi_dot)          \
   _(aten, linalg_norm)               \
   _(aten, linalg_vector_norm)        \

aten/src/ATen/native/BatchLinearAlgebra.cpp

@@ -3836,6 +3836,14 @@ TransposeType to_transpose_type(const bool contig, const bool conj) {
 }
 } // end of anonymous namespace
 
+Tensor& linalg_vecdot_out(const Tensor& x, const Tensor& y, int64_t dim, Tensor& out) {
+  return at::sum_out(out, x * y, /*dim=*/dim);
+}
+
+Tensor linalg_vecdot(const Tensor& x, const Tensor& y, int64_t dim) {
+  return (x * y).sum(/*dim=*/dim);
+}
+
 /*
 Solves the matrix equation AX = B for A triangular.
 'left' If true solves AX = B, if false solves XA = B

aten/src/ATen/native/native_functions.yaml

@@ -10844,6 +10844,13 @@
 - func: linalg_matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
   python_module: linalg
 
+- func: linalg_vecdot(Tensor x, Tensor y, *, int dim=-1) -> Tensor
+  python_module: linalg
+  variants: function
+
+- func: linalg_vecdot.out(Tensor x, Tensor y, *, int dim=-1, Tensor(a!) out) -> Tensor(a!)
+  python_module: linalg
+
 - func: linalg_matrix_exp(Tensor self) -> Tensor
   python_module: linalg
   variants: function

docs/source/linalg.rst

@@ -80,6 +80,7 @@ Matrix Products
 
     cross
     matmul
+    vecdot
     multi_dot
     householder_product
 

test/test_linalg.py

@@ -26,7 +26,7 @@
 from torch.testing import make_tensor
 from torch.testing._internal.common_dtype import (
     all_types, floating_types, floating_and_complex_types, get_all_dtypes, get_all_int_dtypes, get_all_complex_dtypes,
-    get_all_fp_dtypes,
+    get_all_fp_dtypes, all_types_and_complex_and
 )
 from torch.testing._internal.common_cuda import SM53OrLater, tf32_on_and_off, CUDA11OrLater, CUDA9
 from torch.distributions.binomial import Binomial

torch/_torch_docs.py

@@ -3412,6 +3412,10 @@ def merge_dicts(*dicts):
     Unlike NumPy's dot, torch.dot intentionally only supports computing the dot product
     of two 1D tensors with the same number of elements.
 
+.. seealso::
+
+        :func:`torch.linalg.vecdot` the vector product of two batches of vectors along a dimension.
+
 Args:
     input (Tensor): first tensor in the dot product, must be 1D.
     other (Tensor): second tensor in the dot product, must be 1D.

torch/linalg/__init__.py

@@ -2251,3 +2251,42 @@
     >>> torch.dist(Q.mT @ Q, torch.eye(4))
     tensor(6.2158e-07)
 """)
+
+vecdot = _add_docstr(_linalg.linalg_vecdot, r"""
+linalg.vecdot(x, y, *, dim=-1, out=None) -> Tensor
+
+Computes the vector product of two batches of vectors along a dimension.
+
+In symbols, this function computes
+
+.. math::
+
+    \sum_{i=1}^n x_iy_i.
+
+over the dimension :attr:`dim`.
+
+Supports input of half, bfloat16, float, double, cfloat, cdouble and integral dtypes.
+It also supports broadcasting.
+
+.. seealso::
+
+        :func:`torch.matmul` computes a general matrix-matrix multiplication for batches
+        of matrices.
+
+Args:
+    x (Tensor): first batch of vectors.
+    y (Tensor): second batch of vectors.
+
+Keyword args:
+    dim (int): Dimension along which to compute the vector product.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> v1 = torch.randn(3, 2)
+    >>> v2 = torch.randn(3, 2)
+    >>> linalg.vecdot(v1, v2)
+    tensor([ 0.3223,  0.2815, -0.1944])
+    >>> torch.dot(v1[0], v2[0])
+    tensor(0.3223)
+""")

torch/overrides.py

@@ -869,6 +869,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]:
         torch.ravel: lambda input: -1,
         torch.real: lambda input, out=None: -1,
         torch.vdot: lambda input, other, out=None: -1,
+        torch.linalg.vecdot: lambda input, other, dim=-1, out=None: -1,
         torch.view_as_real: lambda input: -1,
         torch.view_as_complex: lambda input: -1,
         torch.reciprocal: lambda input, out=None: -1,

torch/testing/_internal/common_methods_invocations.py

@@ -2012,6 +2012,15 @@ def sample_inputs_isclose(
         yield SampleInput(lhs, args=(rhs,),
                           kwargs=dict(op_kwargs, rtol=rtol, atol=atol, equal_nan=equal_nan))
 
+def sample_inputs_linalg_vecdot(op_info, device, dtype, requires_grad, **kwargs):
+    yield from sample_inputs_binary_pwise(op_info, device, dtype, requires_grad)
+
+    # Add also samples with dim != -1
+    for s in sample_inputs_binary_pwise(op_info, device, dtype, requires_grad):
+        if s.input.ndim > 1:
+            s.kwargs["dim"] = 0
+            yield s
+
 def sample_inputs_t(op_info, device, dtype, requires_grad, **kwargs):
     make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
     return (SampleInput(make_arg((1, 2))),
@@ -9778,6 +9787,17 @@ def ref_pairwise_distance(input1, input2):
            gradcheck_wrapper=gradcheck_wrapper_hermitian_input,
            decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCUDAIfRocm, skipCPUIfNoLapack],
            ),
+    BinaryUfuncInfo('linalg.vecdot',
+                    aten_name='linalg_vecdot',
+                    ref=lambda x, y, *, dim=-1: (x * y).sum(dim),
+                    dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool),
+                    sample_inputs_func=sample_inputs_linalg_vecdot,
+                    supports_forward_ad=True,
+                    supports_fwgrad_bwgrad=True,
+                    skips=(
+                        # torch.sum(out=) has an incorrect behaviour
+                        DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'),
+                    ),),
     OpInfo('linalg.cond',
            aten_name='linalg_cond',
            dtypes=floating_and_complex_types(),