Skip to content

Commit

Permalink
MAINT: simplify axis check in array_api.vecdot
Browse files Browse the repository at this point in the history
  • Loading branch information
seberg committed Aug 1, 2022
1 parent 285084a commit 9d61d5e
Showing 1 changed file with 3 additions and 5 deletions.
8 changes: 3 additions & 5 deletions numpy/array_api/linalg.py
Expand Up @@ -379,16 +379,14 @@ def trace(x: Array, /, *, offset: int = 0) -> Array:
def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
raise TypeError('Only numeric dtypes are allowed in vecdot')
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_ = np.broadcast_arrays(x1._array, x2._array)
x1_ = np.moveaxis(x1_, axis, -1)
x2_ = np.moveaxis(x2_, axis, -1)

if x1_shape[-1] != x2_shape[-1]:
raise ValueError("x1 and x2 must have the same size along the given axis")

res = x1_[..., None, :] @ x2_[..., None]
return Array._new(res[..., 0, 0])

Expand Down

0 comments on commit 9d61d5e

Please sign in to comment.