Vedant Puri

@YichengDWu did this solve the GPU issue?

Thanks for letting me know @YichengDWu. The change works with Lux.jl, but I'm still getting scalar indexing errors on the GPU.

with the change in this PR, this code is working: ```julia using Random using Lux, CUDA, LuxCUDA, ComponentArrays using Zygote, ForwardDiff CUDA.allowscalar(false) #==========================# function testhessian( NN::Lux.AbstractExplicitLayer, data::Tuple; device = cpu_device(),...

But `Lux.Embedding` is failing. This looks like an NNlib issue with ForwardDiff.Dual ```julia using Random using Lux, CUDA, LuxCUDA, ComponentArrays using Zygote, ForwardDiff CUDA.allowscalar(false) #==========================# function testhessian( NN::Lux.AbstractExplicitLayer, data::Tuple; device...

I narrowed the error in the above case to the MWE in https://github.com/FluxML/NNlib.jl/issues/547. That doesn't seem related to this PR.

@abillscmu if you are still interested. @ViralBShah is support for distributed computing with Sundials within scope?

There are some enzyme related errors in NNlib integration tests but they seem unrelated to this PR.

@ToucheSir, LMK I need to add more tests. here's a working MWE with Lux. This also resolves https://github.com/FluxML/Zygote.jl/issues/1348 with the change in this PR, this code is working: ```julia using...

rerunning CI

@stevengj, @devmotion could you do this? > Sorry to just bump this, but I've also run into #899 and wanted to check whether this can be solved. As [JuliaMath/FFTW.jl#249](https://github.com/JuliaMath/FFTW.jl/pull/249) has...