ForwardDiff scalar indexing with GPU arrays

[ChrisRackauckas]

First noted in SciML/OrdinaryDiffEq.jl#2567, but it never got addressed. The same downstream test case used with SparseDiffTools highlights this:

using OrdinaryDiffEq, CUDA, LinearAlgebra, Test, StaticArrays, ADTypes
function f(du, u, p, t)
    mul!(du, A, u)
end
CUDA.allowscalar(false)
A = cu(-rand(3, 3))
u0 = cu([1.0; 0.0; 0.0])
tspan = (0.0f0, 100.0f0)

prob = ODEProblem(f, u0, tspan)
sol = solve(prob, Tsit5()) # works
solve(prob, Rosenbrock23(autodiff = AutoForwardDiff()));

It can probably be minimized directly to the DI call, but I can't do that right now as I don't have the hardware on-hand to debug offline.

Error message: DIgpuerror.txt

But as noted in https://discourse.julialang.org/t/combining-sparsedifftools-jl-and-ordinarydiffeq-jl/130555/9, it's simply because DI uses ForwardDiff.jacobian which is not GPU compatible because it scalar indexes the partials, and it needs to adopt the SparseDiffTools style of partial filling if it wants to be able to vectorize that operation.

[gdalle]

First, here's a non-CUDA MWE leveraging JLArrays, so that people can test this on every computer:

using OrdinaryDiffEq, JLArrays, LinearAlgebra, ADTypes

function f(du, u, p, t)
    A = jl(-ones(3, 3))
    return mul!(du, A, u)
end

u0 = jl([1.0; 0.0; 0.0])
tspan = (0.0f0, 100.0f0)
prob = ODEProblem(f, u0, tspan)
solve(prob, Rosenbrock23(; autodiff=AutoForwardDiff()));

Stack trace:

ulia> solve(prob, Rosenbrock23(; autodiff=AutoForwardDiff()));
ERROR: First call to automatic differentiation for the Jacobian
failed. This means that the user `f` function is not compatible
with automatic differentiation. Methods to fix this include:

1. Turn off automatic differentiation (e.g. Rosenbrock23() becomes
   Rosenbrock23(autodiff = AutoFiniteDiff())). More details can befound at
   https://docs.sciml.ai/DiffEqDocs/stable/features/performance_overloads/
2. Improving the compatibility of `f` with ForwardDiff.jl automatic
   differentiation (using tools like PreallocationTools.jl). More details
   can be found at https://docs.sciml.ai/DiffEqDocs/stable/basics/faq/#Autodifferentiation-and-Dual-Numbers
3. Defining analytical Jacobians. More details can be
   found at https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction

Note: turning off automatic differentiation tends to have a very minimal
performance impact (for this use case, because it's forward mode for a
square Jacobian. This is different from optimization gradient scenarios).
However, one should be careful as some methods are more sensitive to
accurate gradients than others. Specifically, Rodas methods like `Rodas4`
and `Rodas5P` require accurate Jacobians in order to have good convergence,
while many other methods like BDF (`QNDF`, `FBDF`), SDIRK (`KenCarp4`),
and Rosenbrock-W (`Rosenbrock23`) do not. Thus if using an algorithm which
is sensitive to autodiff and solving at a low tolerance, please change the
algorithm as well.

Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] jacobian!(J::JLArray{…}, f::Function, x::JLArray{…}, fx::JLArray{…}, integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, jac_config::Tuple{…})
    @ OrdinaryDiffEqDifferentiation ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_wrappers.jl:223
  [2] calc_J!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:222 [inlined]
  [3] calc_W!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:627 [inlined]
  [4] calc_W!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:565 [inlined]
  [5] calc_rosenbrock_differentiation!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:702 [inlined]
  [6] perform_step!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqRosenbrock.Rosenbrock23Cache{…}, repeat_step::Bool)
    @ OrdinaryDiffEqRosenbrock ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:46
  [7] perform_step!
    @ ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:27 [inlined]
  [8] solve!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:620
  [9] __solve(::ODEProblem{…}, ::Rosenbrock23{…}; kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:7
 [10] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:1 [inlined]
 [11] #solve_call#36
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:670 [inlined]
 [12] solve_call
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:627 [inlined]
 [13] #solve_up#45
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1224 [inlined]
 [14] solve_up
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1201 [inlined]
 [15] #solve#43
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1096 [inlined]
 [16] solve(prob::ODEProblem{…}, args::Rosenbrock23{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1086
 [17] top-level scope
    @ ~/Documents/GitHub/Julia/_Scratchpad2/mwe.jl:1

caused by: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:50 [inlined]
  [6] seed!(duals::JLArray{…}, x::JLArray{…}, seeds::Tuple{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/Wq9Wb/src/apiutils.jl:84
  [7] vector_mode_dual_eval!
    @ ~/.julia/packages/ForwardDiff/Wq9Wb/src/apiutils.jl:29 [inlined]
  [8] vector_mode_jacobian!(result::JLArray{…}, f!::SciMLBase.UJacobianWrapper{…}, y::JLArray{…}, x::JLArray{…}, cfg::ForwardDiff.JacobianConfig{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/Wq9Wb/src/jacobian.jl:157
  [9] jacobian!
    @ ~/.julia/packages/ForwardDiff/Wq9Wb/src/jacobian.jl:82 [inlined]
 [10] jacobian!(::SciMLBase.UJacobianWrapper{…}, ::JLArray{…}, ::JLArray{…}, ::DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgJacobianPrep{…}, ::AutoForwardDiff{…}, ::JLArray{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/sPszY/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl:489
 [11] jacobian!(J::JLArray{…}, f::Function, x::JLArray{…}, fx::JLArray{…}, integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, jac_config::Tuple{…})
    @ OrdinaryDiffEqDifferentiation ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_wrappers.jl:221
 [12] calc_J!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:222 [inlined]
 [13] calc_W!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:627 [inlined]
 [14] calc_W!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:565 [inlined]
 [15] calc_rosenbrock_differentiation!
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/derivative_utils.jl:702 [inlined]
 [16] perform_step!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqRosenbrock.Rosenbrock23Cache{…}, repeat_step::Bool)
    @ OrdinaryDiffEqRosenbrock ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:46
 [17] perform_step!
    @ ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:27 [inlined]
 [18] solve!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:620
 [19] __solve(::ODEProblem{…}, ::Rosenbrock23{…}; kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:7
 [20] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:1 [inlined]
 [21] #solve_call#36
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:670 [inlined]
 [22] solve_call
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:627 [inlined]
 [23] #solve_up#45
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1224 [inlined]
 [24] solve_up
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1201 [inlined]
 [25] #solve#43
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1096 [inlined]
 [26] solve(prob::ODEProblem{…}, args::Rosenbrock23{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1086
 [27] top-level scope
    @ ~/Documents/GitHub/Julia/_Scratchpad2/mwe.jl:1
Some type information was truncated. Use `show(err)` to see complete types.

I'm not sure what you expect from DI here.
This bug is caused by ForwardDiff itself not being compatible with GPUs. Assuming we want to fix it, it should be fixed in ForwardDiff by adding specific dispatches for AbstractGPUArray in an extension depending on GPUArraysCore. DI's purpose is to forward as much as possible to the backend, and interfere only when absolutely necessary.
As a side note, I also don't understand why SparseDiffTools used to solve this problem, since the Jacobian we compute here is a dense one?

[ChrisRackauckas]

OrdinaryDiffEq.jl always used SparseDiffTools on non-static arrays (or rather, in-place functions) because its Jacobian function was just faster. It has a trade-off because it is building the array-of-array of partials and storing that in memory, and then for the Jacobian calls it vectorizes over the partials matrix. This of course was done at first in order to easily represent the color vectors, but it had the upside of getting better parallel performance and GPU support. The downside of course is that you can more easily differentiate the basis vectors by just bumping +-1 on i and i+1, thus using a lot less memory in the "standard" case (though with chunking, you do have to at least have a chunk sized amount created). However, we found that the SparseDiffTools way just had too many upsides, since normally if someone really had memory issues with a Jacobian the answer is to go sparse anyways, and the overhead is <2x in memory while usually being slightly faster, and it had better support for weird array types, so we ended up just always using it for the last half decade.

The question now is that DI basically removes those improvements and we're back to the ForwardDiff issues, and what should we do about that. We could say that ForwardDiff should just get better, but that has moved so slow that it's kind of a non-answer. We could add a AutoSDTForwardDiff ADType and keep that alive, though I'd hope to pass off the maintenance if I could 😅. We could just make this be a separate ForwardDiff jacobian call that is kept in DI, with some way in the AutoForwardDiff to trigger it. Whatever the solution is its up to you, but this has been waiting for awhile so I'd hope to get something in like in a few weeks.

Does DI sparse already do the vectorization for the color vectors?

[gdalle]

it had the upside of getting better parallel performance and GPU support.

I understand that the use case here is a function f(x) whose input is a x::CuArray but what type do you want to have for its Jacobian? Matrix / SparseMatrixCSC / CuMatrix / CuSparseMatrixCSR?

Does DI sparse already do the vectorization for the color vectors?

It has a trade-off because it is building the array-of-array of partials and storing that in memory, and then for the Jacobian calls it vectorizes over the partials matrix.

What do you mean by vectorization? I'm assuming something different than what ForwardDiff does, which is to process 12 basis vectors at a time within a chunk? In sparse mode, DI with ForwardDiff processes 12 multibasis vectors at a time, if that answers your question.

I looked at the code for SparseDiffTools's Jacobian which you linked in adrhill/SparseConnectivityTracer.jl#250 but it is not very enlightening.

We could say that ForwardDiff should just get better, but that has moved so slow that it's kind of a non-answer.

That would definitely be my first choice. Especially since it's about adding functionality in an extension for stuff which currently errors (gradient and jacobian on GPU), I don't see a convincing case against it.

this has been waiting for awhile so I'd hope to get something in like in a few weeks.

From my perspective, this has been waiting for exactly 2 hours, since you opened the issue ;) If you don't report the problems you encounter, no solution can be found. Keep in mind that it's rather hard for outsiders to figure out which tests are supposed to fail and which tests are not, especially in big PRs that last months like SciML/OrdinaryDiffEq.jl#2567, so opening issues never hurts.

[gdalle]

To me it looks like the solution is reviving (something like) JuliaDiff/ForwardDiff.jl#619 and putting it in a GPUArraysCore extension to ForwardDiff so that it brings no regression, only new functionality.

[ChrisRackauckas]

CuMatrix / CuSparseMatrixCSR

What do you mean by vectorization? I'm assuming something different than what ForwardDiff does, which is to process 12 basis vectors at a time within a chunk? In sparse mode, DI with ForwardDiff processes 12 multibasis vectors at a time, if that answers your question.

It's not about the chunking, it's how the chunks are made. Basically this part right here https://github.com/JuliaDiff/SparseDiffTools.jl/blob/master/src/differentiation/compute_jacobian_ad.jl#L41-L46 where if you already have the colorvec, you can define the partials via a broadcast call over i == color. Then using those partial vectors, you can just broadcast them into the dual to set them https://github.com/JuliaDiff/SparseDiffTools.jl/blob/master/src/differentiation/compute_jacobian_ad.jl#L204. So it's all broadcast.

ForwardDiff seeds here https://github.com/JuliaDiff/ForwardDiff.jl/blob/master/src/jacobian.jl#L197 which seeds by loops https://github.com/JuliaDiff/ForwardDiff.jl/blob/d81302372bc812fbc7f3e2eafa17b4b8f570d8ae/src/apiutils.jl#L73-L106. You could do something tricky like what SDT does with broadcasting of getindex with index arrays https://github.com/JuliaDiff/SparseDiffTools.jl/blob/master/src/differentiation/compute_jacobian_ad.jl#L404-L408 in order to change that loop to a broadcast operation.

[ChrisRackauckas]

To me it looks like the solution is reviving (something like) JuliaDiff/ForwardDiff.jl#619 and putting it in a GPUArraysCore extension to ForwardDiff so that it brings no regression, only new functionality.

The only issue I had with that was that keeping most of the ForwardDiff.jacobian structure makes it fairly hard to match/beat performance, because making partial seeds on-demand makes it naturally harder to vectorize vs having the color vectors already there and just broadcasting over them. That PR got bogged down in exactly the performance things I mentioned. While if you also have a system that does coloring, you already have a system for getting the batches of all possible colors upfront. It's just much easier to make the coloring version of the Jacobian be GPU-compatible, and then just use that with the trivial e_i colors for the Jacobian.

[gdalle]

I may have something for you here. DI's fallback jacobian utility is designed to be fully GPU-friendly, using only mapping and broadcasting operations. It also interacts well with preparation, where the basis vectors are precomputed only once and then looped over at execution time. Here's the full code running at execution time.
Normally, that fallback doesn't kick in if the backend provides its own optimized jacobian implementation (as is the case for ForwardDiff). However, DI has a special backend wrapper which forgets every implementation except pushforward or pullback. So the code below uses ForwardDiff's pushforward inside DI's fallback jacobian. As you can see, it works on GPU arrays:

julia> import DifferentiationInterface as DI, ForwardDiff, JLArrays

julia> x = JLArrays.jl(ones(2));

julia> backend = DI.AutoForwardFromPrimitive(DI.AutoForwardDiff());

julia> DI.jacobian(identity, backend, x)
2×2 JLArray{Float64, 2}:
 1.0  0.0
 0.0  1.0

julia> DI.jacobian(identity, DI.AutoForwardDiff(), x)
ERROR: Scalar indexing is disallowed.

This AutoForwardFromPrimitive was meant to be internal, but if it is of any use to anyone (e.g. to bypass scalar indexing in ForwardDiff's seeding) we can make it public. I haven't tested it on CUDA arrays, but I assume it would just return a CuMatrix without complaining, as it is doing for JLArrays.

As for CuSparseMatrixCSR, we already have an optimized implementation, which is activated if the sparsity pattern you provide is also of the same type. You can use it as follows:

using ADTypes
using CUDA, CUDA.CUSPARSE
using DifferentiationInterface
using ForwardDiff: ForwardDiff
using LinearAlgebra
using SparseConnectivityTracer
using SparseMatrixColorings

f(x) = map(abs2, x)

x = CuVector(rand(Float32, 10));

S_cpu = jacobian_sparsity(f, Array(x), TracerSparsityDetector())
S = CuSparseMatrixCSR(Float32.(S_cpu))

sparsity_detector = ADTypes.KnownJacobianSparsityDetector(S)
coloring_algorithm = GreedyColoringAlgorithm()
backend = AutoSparse(AutoForwardDiff(); sparsity_detector, coloring_algorithm);

prep = prepare_jacobian(f, backend, x);

jacobian(f, prep, backend, x)  # returns a CuSparseMatrixCSR

This code is from the notebook we discussed on Discourse. I don't know why there was a warning for scalar indexing in the case of the Hessian, because there shouldn't be any scalar indexing outside of preparation. It seems to be gone for the case of the Jacobian, I'll investigate further.
EDIT: the scalar indexing warning was because I was computing a Hessian with ForwardDiff-over-ForwardDiff, so there was scalar indexing inside the gradient part of the HVP.

[gdalle]

When I try it on your Rosenbrock solver example, it fails at a different point, but I think that one is not my fault. It might be due to the lack of good factorization methods for JLArrays, so possibly the error disappears with CUDA?

julia> solve(
           prob,
           Rosenbrock23(;
               autodiff=DifferentiationInterface.AutoForwardFromPrimitive(AutoForwardDiff())
           ),
       );
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:50 [inlined]
  [6] scalar_getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:36 [inlined]
  [7] _getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:19 [inlined]
  [8] getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:17 [inlined]
  [9] getindex
    @ ./subarray.jl:320 [inlined]
 [10] _unsafe_getindex_rs
    @ ./reshapedarray.jl:276 [inlined]
 [11] _unsafe_getindex
    @ ./reshapedarray.jl:273 [inlined]
 [12] getindex
    @ ./reshapedarray.jl:261 [inlined]
 [13] getindex
    @ ./subarray.jl:320 [inlined]
 [14] iterate
    @ ./iterators.jl:137 [inlined]
 [15] iterate
    @ ./iterators.jl:134 [inlined]
 [16] _any(f::typeof(iszero), itr::Base.Iterators.Reverse{SubArray{…}}, ::Colon)
    @ Base ./reduce.jl:1236
 [17] any
    @ ./reduce.jl:1228 [inlined]
 [18] _notsuccessful
    @ ~/.julia/packages/LinearSolve/wrTJy/src/LinearSolve.jl:172 [inlined]
 [19] macro expansion
    @ ~/.julia/packages/LinearSolve/wrTJy/src/factorization.jl:9 [inlined]
 [20] solve!(cache::LinearSolve.LinearCache{…}, alg::LinearSolve.QRFactorization{…}; kwargs::@Kwargs{…})
    @ LinearSolve ~/.julia/packages/LinearSolve/wrTJy/src/factorization.jl:1
 [21] solve!
    @ ~/.julia/packages/LinearSolve/wrTJy/src/factorization.jl:1 [inlined]
 [22] macro expansion
    @ ~/.julia/packages/LinearSolve/wrTJy/src/default.jl:399 [inlined]
 [23] solve!(::LinearSolve.LinearCache{…}, ::LinearSolve.DefaultLinearSolver; assump::LinearSolve.OperatorAssumptions{…}, kwargs::@Kwargs{…})
    @ LinearSolve ~/.julia/packages/LinearSolve/wrTJy/src/default.jl:354
 [24] solve!
    @ ~/.julia/packages/LinearSolve/wrTJy/src/default.jl:354 [inlined]
 [25] #solve!#11
    @ ~/.julia/packages/LinearSolve/wrTJy/src/common.jl:299 [inlined]
 [26] solve!
    @ ~/.julia/packages/LinearSolve/wrTJy/src/common.jl:298 [inlined]
 [27] #dolinsolve#2
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/linsolve_utils.jl:29 [inlined]
 [28] dolinsolve
    @ ~/.julia/packages/OrdinaryDiffEqDifferentiation/7QyLE/src/linsolve_utils.jl:5 [inlined]
 [29] perform_step!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqRosenbrock.Rosenbrock23Cache{…}, repeat_step::Bool)
    @ OrdinaryDiffEqRosenbrock ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:58
 [30] perform_step!
    @ ~/.julia/packages/OrdinaryDiffEqRosenbrock/1cjFj/src/rosenbrock_perform_step.jl:27 [inlined]
 [31] solve!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:620
 [32] __solve(::ODEProblem{…}, ::Rosenbrock23{…}; kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:7
 [33] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/zs1s7/src/solve.jl:1 [inlined]
 [34] #solve_call#36
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:670 [inlined]
 [35] solve_call
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:627 [inlined]
 [36] #solve_up#45
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1224 [inlined]
 [37] solve_up
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1201 [inlined]
 [38] #solve#43
    @ ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1096 [inlined]
 [39] solve(prob::ODEProblem{…}, args::Rosenbrock23{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/yshbt/src/solve.jl:1086
 [40] top-level scope
    @ ~/Documents/GitHub/Julia/_Scratchpad2/mwe.jl:1
Some type information was truncated. Use `show(err)` to see complete types.

[ChrisRackauckas]

When I try it on your Rosenbrock solver example, it fails at a different point, but I think that one is not my fault. It might be due to the lack of good factorization methods for JLArrays, so possibly the error disappears with CUDA?

Yeah, that failure is because there's no LU/QR factorization defined for JLArray. So that means the Jacobian building worked and CUDA should be fine. So okay, it seems like you have a way that is GPU compatible, it's just not called because ForwardDiff.jacobian is preferred. Is this something we should make a switch for in the ADType? Or specialize on CuMatrix / ArrayInterface? Generally the way to do this is to hit that other code if ArrayInterface.fast_scalar_indexing(J) == true

[gdalle]

Is this something we should make a switch for in the ADType?

As explained above, DI already has a built-in way to deactivate custom implementations for everything except pushforward / pullback. If it is useful outside of internal testing, I can make AutoForwardFromPrimitive public to bypass ForwardDiff.jacobian. Then it's up to the users to figure out which one to pick.

[ChrisRackauckas]

Okay yeah that would be very useful to expose for this, and if !ArrayInterface.fast_scalar_indexing(J) it would make sense to default to it.

[gdalle]

DI currently does not depend on ArrayInterface, and I want to keep it that way. Besides, it's not my job to decide what ForwardDiff does with its Jacobians. If users find that ForwardDiff doesn't work for a specific case (like GPU arrays), then it's up to them to bypass its higher-level constructs through AutoForwardFromPrimitive

[ChrisRackauckas]

Well we'd need to make it work right at least for the SciML uses. We could just add a SciMLDIOverloads.jl that overloads the jacobian function to check the array interface compat and then send it to the right dispatch. What's the right level to overload that?

[gdalle]

Would it be enough to specialize the dispatch on GPU arrays? Or are there other cases where ForwardDiff fails?

[ChrisRackauckas]

What about other cases where you want to not hit a sclarizing dispatch like ArrayPartition? What's the reason to ignore traits and make everyone write an extension for DI? It would be weird if DI is the only thing in the ecosystem that requires special casing while everything else just uses high level traits.