import warp
import numpy as np


@warp.kernel
def kernel(
    a: warp.array(dtype=bool),
    b: warp.array(dtype=float),
):
    data = warp.tile_load(a, TILE_SIZE)
    data0 = warp.untile(data)
    data1 = warp.where(data0, 1., 0.)
    data2 = warp.tile(data1)
    warp.tile_store(b, data2)


m = 128
TILE_SIZE = m


data_a = warp.array(np.random.rand(m) < 0.3, dtype=bool)
data_b = warp.array(np.zeros(m), dtype=float)

warp.launch_tiled(kernel, dim=[1], inputs=[data_a, data_b], block_dim=128)

print(data_b.numpy())