Fixup more than simple jacobian

[wsmoses]

No description provided.

[wsmoses]

@swilliamson7 does this fix your jacobian case?

[github-actions]

Benchmark Results

	main	ad01cc5...	main/ad01cc558b5d4f...
basics/overhead	4.03 ± 0.001 ns	4.03 ± 0.01 ns	1
time_to_load	0.456 ± 0.00077 s	0.454 ± 0.0092 s	1.01

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

[swilliamson7]

@swilliamson7 does this fix your jacobian case?

What does the variable n_out_val mean exactly? It looks like I need to provide it.

julia> grad = jacobian(Reverse, compute_jacobian, S)
ERROR: MethodError: no method matching jacobian(::ReverseMode{false, FFIABI, false, false}, ::typeof(compute_jacobian), ::ShallowWaters.ModelSetup{Float32, Float32})

Closest candidates are:
  jacobian(::ReverseMode{ReturnPrimal, RABI, ErrIfFuncWritten}, ::F, ::X, ::Val{n_out_val}, ::Val{1}) where {F, X, n_out_val, ReturnPrimal, RABI<:EnzymeCore.ABI, ErrIfFuncWritten}
   @ Enzyme ~/.julia/packages/Enzyme/afdHl/src/Enzyme.jl:1302
  jacobian(::ReverseMode{ReturnPrimal, RABI, ErrIfFuncWritten}, ::F, ::X, ::Val{n_out_val}) where {F, X, n_out_val, ReturnPrimal, RABI<:EnzymeCore.ABI, ErrIfFuncWritten}
   @ Enzyme ~/.julia/packages/Enzyme/afdHl/src/Enzyme.jl:1302
  jacobian(::ReverseMode{ReturnPrimal, RABI, ErrIfFuncWritten}, ::F, ::X, ::Val{n_out_val}, ::Val{chunk}) where {F, X, chunk, n_out_val, ReturnPrimal, RABI<:EnzymeCore.ABI, ErrIfFuncWritten}
   @ Enzyme ~/.julia/packages/Enzyme/afdHl/src/Enzyme.jl:1253
  ...

Stacktrace:
 [1] top-level scope
   @ REPL[8]:1

[wsmoses]

it's the number of outputs

[swilliamson7]

If f doesn't necessarily return an array but rather a structure, what would I put? Or should I define f to return an array (say my state variable)?

[wsmoses]

@swilliamson7 just pushed a version which computes n_outs for you.

Currently reverse mode jacobian allows any input type, but requires an array output type.
Currently forward mode jacobian allows any output type, but requires an array or scalar input type.

[swilliamson7]

@swilliamson7 just pushed a version which computes n_outs for you.

Currently reverse mode jacobian allows any input type, but requires an array output type. Currently forward mode jacobian allows any output type, but requires an array or scalar input type.

I saw! That was speedy lol, it doesn't seem like I can test it yet, or at least updating Enzyme doesn't bring any changes. I'm guessing I need to wait for the checks to pass or something?

[wsmoses]

@swilliamson7 it was just merged to main, so you now should be able to just do add Enzyme#main

[swilliamson7]

I got a fun blank-ish error,

julia> grad = jacobian(Reverse, compute_jacobian, S)
ERROR: MethodError: no method matching (ReverseMode{false, FFIABI, false})()
Stacktrace:
 [1] jacobian(::ReverseMode{false, FFIABI, false, false}, f::typeof(compute_jacobian), x::ShallowWaters.ModelSetup{Float32, Float32})
   @ Enzyme ~/.julia/packages/Enzyme/BkuKz/src/Enzyme.jl:1500
 [2] top-level scope
   @ REPL[13]:1