def METHOD_NAME(dims, axis, iava, ncp):
    """Compute restriction mask when using cupy arrays"""
    otherdims = np.array(dims)
    otherdims = np.delete(otherdims, axis)
    iavamask = ncp.zeros(int(dims[axis]), dtype=int)
    iavamask[iava] = 1
    iavamask = ncp.moveaxis(
        ncp.broadcast_to(iavamask, list(otherdims) + [dims[axis]]), -1, axis
    )
    iavamask = ncp.where(iavamask.ravel() == 1)[0]
    return iavamask