How to assign values for ndarray via indexing


mask = a < b
a[mask] = b[mask]

The above assignment does NOT work (both a and b are in shape [B,N,C]). What is the correct way to assign values for this? Thanks


The output of the mask in numpy is a boolean array. NDArray doesn’t support boolean dtype. Another way in numpy is to use np.where(condition, [x, y]) so the assignment will be a = np.where(a < b, b, a). Equivalently for NDArray, you can use a = nd.where(a < b, b, a). Note that, for inplace update, you need to wither use a[:] or nd.where(a < b, b, a, out=a) as shown here.