import warp as wp

@wp.kernel
def foo(arr:wp.array(dtype=wp.int32, ndim=2)):
    i, j = wp.tid()
    arr[i, j] = i * j

with wp.ScopedDevice("cpu"):
    arr = wp.zeros(shape=(3, 3), strides=(4, 12), dtype=wp.int32)
    wp.launch(foo, dim=arr.shape, inputs=(arr,))
    print(arr)