Update the torch vecdot wrapper

- torch.vecdot incorrectly allows broadcasting in the contracted dimensions
- The 2023 version of the spec updates the language to require axis to apply
  before broadcasting, not after. The implementation for integer arguments is
  updated to follow this behavior. Note that the spec only actually requires
  axis to be negative, but we allow nonnegative axis too if x1 and x2 have the
  required number of dimensions and those dimensions have same value, which
  matches the numpy gufunc vecdot implementation and also the base
  torch.vecdot implementation as far as I can tell.

Author: asmeurer

array_api_compat/torch/linalg.py

@@ -40,19 +40,18 @@ def vecdot(x1: array, x2: array, /, *, axis: int = -1, **kwargs) -> array:
 
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
 
+    # torch.linalg.vecdot incorrectly allows broadcasting along the contracted dimension
+    if x1.shape[axis] != x2.shape[axis]:
+        raise ValueError("x1 and x2 must have the same size along the given axis")
+
     # torch.linalg.vecdot doesn't support integer dtypes
     if isdtype(x1.dtype, 'integral') or isdtype(x2.dtype, 'integral'):
         if kwargs:
             raise RuntimeError("vecdot kwargs not supported for integral dtypes")
-        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)
+
+        x1_ = torch.moveaxis(x1, axis, -1)
+        x2_ = torch.moveaxis(x2, axis, -1)
+        x1_, x2_ = torch.broadcast_tensors(x1_, x2_)
 
         res = x1_[..., None, :] @ x2_[..., None]
         return res[..., 0, 0]