Add implementation for torch.linalg.vecdot for integer dtypes

Author: asmeurer

array_api_compat/torch/linalg.py

@@ -22,6 +22,27 @@ def cross(x1: array, x2: array, /, *, axis: int = -1) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
     return torch_linalg.cross(x1, x2, dim=axis)
 
-__all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot']
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
+    from ._aliases import isdtype
+
+    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
+
+    # torch.linalg.vecdot doesn't support integer dtypes
+    if isdtype(x1.dtype, 'integral') or isdtype(x2.dtype, 'integral'):
+        ndim = max(x1.ndim, x2.ndim)
+        x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
+        x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
+        if x1_shape[axis] != x2_shape[axis]:
+            raise ValueError("x1 and x2 must have the same size along the given axis")
+
+        x1_, x2_ = torch.broadcast_tensors(x1, x2)
+        x1_ = torch.moveaxis(x1_, axis, -1)
+        x2_ = torch.moveaxis(x2_, axis, -1)
+
+        res = x1_[..., None, :] @ x2_[..., None]
+        return res[..., 0, 0]
+    return torch.linalg.vecdot(x1, x2, axis=axis)
+
+__all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot', 'vecdot']
 
 del linalg_all