Adding GTPSA.jl to the interface

[mattsignorelli]

I'll use this issue for questions I have for implementing this interface into GTPSA.jl. I've done some reading of the documentation and those AD packages in ext, and here are some of my first questions

1. As long as I also overload value_and_gradient, value_and_jacobian, etc, is it safe to assume pushforward, value_and_pushforward only receive single variable functions?
2. Should I basically plan to overload all of the functions here?

[gdalle]

Thanks for contributing!
The documentation on writing extensions is basically non-existent, sorry for that. The reason is that I wrote all of the extensions myself, and they're not the kind of thing many people will try to touch. I'll write some more docs today and ping you.

[adrhill]

As long as I also overload value_and_gradient, value_and_jacobian, etc, is it safe to assume pushforward, value_and_pushforward only receive single variable functions?

For forward-mode AD, it's technically enough to only implement pushforward.
However, we recommend calling backend operators whenever possible:

While it is possible to define every operator using just pushforward and pullback, some backends have more efficient implementations of high-level operators. When they are available, we always call these backend-specific overloads.

Traits are used to indicate support for two-argument functions f!(y, x)

You might want to take a look at how overloads are handled on other backends: https://docs.juliahub.com/General/DifferentiationInterface/stable/overloads/

[gdalle]

I have written a detailed dev guide in #317

[gdalle]

You can now read it here: https://gdalle.github.io/DifferentiationInterface.jl/DifferentiationInterface/dev/dev_guide/

[mattsignorelli]

So are the arguments x, dx, y, and dy, in pushforward always mutable structs (vectors), or can they also be scalars? I guess I need to check if the output of f is also a vector or scalar in pushforward? So there's several checks for if x isa Number etc?

[gdalle]

DifferentiationInterface supports functions that take either a scalar or an array as input, and return either a scalar or an array. For some operators, like the Jacobian, there's only one possible combination (array to array). But for the pushforward, there are 4 different combinations, which you must all handle. This can be done either with checks inside the function or with multiple dispatch.

https://gdalle.github.io/DifferentiationInterface.jl/DifferentiationInterface/dev/operators/#List-of-operators

[gdalle]

Note that you don't need to define the mutable version pushforward! when the output y is a number, because dy won't be mutable

[mattsignorelli]

For pushforward where y=f(x) is some vector result, presumably it should return the Jacobian. However I am having trouble using similar to do so, since both x and y are vectors. If someone is using StaticArrays for example, then to produce a SizedMatrix result I would need to do

similar(y, GTPSA.numtype(eltype(y)), Size(length(y),length(x))

But Size is in StaticArrays, not in Base, and I see StaticArrays isn't a dependency of DifferentiationInterface. How should I handle this?

[gdalle]

For pushforward where y=f(x) is some vector result, presumably it should return the Jacobian.

No, the pushforward is the product of the Jacobian with a given vector, also called JVP. The whole point is to avoid materializing the full matrix if you can.

If someone is using StaticArrays for example

Each operator has an in-place and an out-of-place version, so we only expect the out-of-place version to work with StaticArrays. The test package DifferentiationInterfaceTest contains a list of static_scenarios() with which you can test a backend implementation for StaticArrays compatibility. But indeed, we don't want StaticArrays as a dependency.

Perhaps I could lend a hand if you point me to your code draft?

[mattsignorelli]

No, the pushforward is the product of the Jacobian with a given vector, also called JVP. The whole point is to avoid materializing the full matrix if you can.

Ah I see now, my understanding was completely wrong

Perhaps I could lend a hand if you point me to your code draft?

Sure! Here is the fork I'm working on. I have derivative, second_derivate, gradient, jacobian, hessian all implemented for onearg

[gdalle]

I strongly recommend adding a test file as outlined in the docs, and running the Single/GPTSA tests regularly. That's how you're gonna catch the bugs

[mattsignorelli]

Yes I will soon, I just haven't had enough time to work on this yet

[mattsignorelli]

OK I am having a bit of trouble and could use some help. I'm not the most familiar with how to work with extensions.

First I tried defining the AutoGTPSA type in the module in ext, but then ran into problems because exporting AutoGTPSA from DifferentiationInterface leaves AutoGTPSA undefined and weird things happened

Then I also dev-ed out ADTypes, and defined AutoGTPSA there following the other backends. This solved the first problem but now I still cannot load AutoGTPSA:

If you dev my DifferentiationInterface fork (note branch is called addgtpsa) and my ADTypes fork (branch also called addgtpsa), and then in Julia do:

julia> using DifferentiationInterface

julia> import GTPSA

julia> DifferentiationInterface.check_available(AutoGTPSA())
false

Any help is much appreciated!

[gdalle]

Can you open a PR to the present repo so that I may modify it?

[mattsignorelli]

PR opened

Here's the definition for the AutoGTPSA type:

"""
    AutoGTPSA{D}

Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation.

# Constructors

    AutoGTPSA(; descriptor=nothing)

# Fields

  - `descriptor::D`: can be either

      + a GTPSA `Descriptor` specifying the number of variables/parameters, parameter 
        order, individual variable/parameter truncation orders, and maximum order. See 
        [here](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) for more details.
      + `nothing` to automatically use a `Descriptor` given the context
      
"""
Base.@kwdef struct AutoGTPSA{D} <: AbstractADType 
  descriptor::D = nothing
end

mode(::AutoGTPSA) = ForwardMode()