@wp.kernel
def test_tile_assign_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    i, j = wp.tid()

    a = wp.tile_zeros(shape=(8,), dtype=float)

    a[j] = x[j]

    wp.tile_atomic_add(y, a, offset=(0,))

x = wp.full(8, 2.0, dtype=float, device=device, requires_grad=True)
y = wp.zeros(8, dtype=float, device=device, requires_grad=True)

wp.launch(test_tile_assign_kernel, dim=[1, 8], inputs=[x], outputs=[y], block_dim=64, device=device)