import numpy as np

rng = np.random.default_rng(seed=42)

a = rng.standard_normal((100, 20))
b = rng.standard_normal((100,))

_, n = a.shape
x = np.zeros(n)

w = np.linalg.vecdot(b - a @ x, a, axis=0)  # works

import numpy as np
import jax.numpy as jnp

rng = np.random.default_rng(seed=42)

a = rng.standard_normal((100, 20))
b = rng.standard_normal((100,))

a_j = jnp.array(a)
b_j = jnp.array(b)

_, n_j = a_j.shape
x_j = jnp.zeros(n_j)

w_j = jnp.linalg.vecdot(b_j - a_j @ x_j, a_j, axis=0)  # works

import numpy as np
import cupy

rng = np.random.default_rng(seed=42)

a = rng.standard_normal((100, 20))
b = rng.standard_normal((100,))

a_cu = cupy.asarray(a)
b_cu = cupy.asarray(b)

_, n_cu = a_cu.shape
x_cu = cupy.zeros(n_cu)

w_cu = cupy.dot(b_cu - a_cu @ x_cu, a_cu)  # works

import numpy as np
import cupy.array_api as cp

rng = np.random.default_rng(seed=42)

a = rng.standard_normal((100, 20))
b = rng.standard_normal((100,))

a_c = cp.asarray(a)
b_c = cp.asarray(b)

_, n_c = a_c.shape
x_c = cp.zeros(n_c)

w_c = cp.linalg.vecdot(b_c - a_c @ x_c, a_c, axis=0)  # does not work
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[48], line 1
----> 1 w_c = cp.linalg.vecdot(b_c - a_c @ x_c, a_c, axis=0)

File /mnt/c/Users/Saransh/Saransh_softwares/OpenSource/Python/glass/.env/lib/python3.10/site-packages/cupy/array_api/linalg.py:385, in vecdot(x1, x2, axis) 
    383 x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
    384 if x1_shape[axis] != x2_shape[axis]:
--> 385     raise ValueError("x1 and x2 must have the same size along the given axis")
    387 x1_, x2_ = np.broadcast_arrays(x1._array, x2._array)
    388 x1_ = np.moveaxis(x1_, axis, -1)

ValueError: x1 and x2 must have the same size along the given axis