using Flux: Flux
using Enzyme: Enzyme
using CUDA: CUDA
using cuDNN: cuDNN

function reproducer(; device = Flux.gpu)
    Enzyme.API.strictAliasing!(false)
    Enzyme.Compiler.VERBOSE_ERRORS[] = true
    model =
        Flux.Chain(Flux.Dense(28^2 => 32, Flux.sigmoid), Flux.Dense(32 => 10), Flux.softmax) |>
        device
    # this allocates space for the gradient
    model_shadow = Enzyme.Duplicated(model) |> device
    input = Flux.randn32(28^2, 1) |> device
    label = [i == 3 for i in 1:10]
    f = (m, x, y) -> sum(abs2, m(x) .- y)
    grad_model = Flux.gradient(
        f,
        model_shadow,
        # input and label are const because we differentiate w.r.t. to the model parameters, not input or output.
        Enzyme.Const(input),
        Enzyme.Const(label),
    )
end