from typing import Any
import warp as wp

wp.init()
wp.config.verify_autograd_array_access = True


@wp.func
def read_func(x: wp.array(dtype=Any), tid: int):
    return x[tid]


@wp.kernel
def safe_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    tid = wp.tid()
    a = read_func(x, tid)
    y[tid] = a


@wp.kernel
def problem_kernel(x: wp.array(dtype=float)):
    tid = wp.tid()
    a = read_func(x, tid)
    x[tid] = a + 1.0


device = "cuda:0"

N = 3
x = wp.ones(N, dtype=float, requires_grad=True)
y = wp.zeros(N, dtype=float, requires_grad=True)

wp.launch(safe_kernel, N, inputs=[x, y], device=device)

x = wp.ones(N, dtype=float, requires_grad=True)

wp.launch(problem_kernel, N, inputs=[x], device=device)