Using DifferentiationInterface in SymbolicRegression

[MilesCranmer]

Hi @gdalle,

Awesome package, thanks for making it!

I am considering whether to use this as the backend interface of SymbolicRegression.jl to allow the user to pick an AD backend (current attempt). However, one thing I was wondering if you would consider is designating a mapping from Symbol to the correct AbstractADType? I could implement this within my package, but it seems like something that should live here for others as well.

Basically, I would prefer that my users do not need to install both Zygote.jl and DifferentiationInterface.jl (or ADTypes) in their environment to set Zygote as the backend. I would rather they just pass, e.g., :Zygote and then have my package, which depends on DifferentiationInterface.jl, be able to internally map to the right AbstractADType.

Within DifferentiationInterface.jl this could look like:

function lookup_backend(key::Symbol)
    return ADTYPE_MAP[key]
end

const ADTYPE_MAP = Dict{Symbol,AbstractADType}([
    :Zygote => AutoZygote(),
    :Enzyme => AutoEnzyme(),
    :ForwardDiff => AutoForwardDiff(),
])

Then within my package, I can have my Options constructor look like this:

using DifferentiationInterface: lookup_backend, AbstractADType

function Options(autodiff_backend::Union{AbstractADType,Symbol})
    backend = autodiff_backend isa Symbol ? lookup_backend(autodiff_backend) : autodiff_backend

    return Options{typeof(backend)}(backend)
end

which means the user has a more lightweight interface; they can just pass the symbol:

using SymbolicRegression, Zygote

options = Options(:Zygote)

And then I map to the right type internally.

I guess I could also re-export all the abstract types but it feels a bit overkill. That would also mean I would be unable to provide helpful error messages if they forget the name of the backend and pass in the wrong symbol.

What do you think?

Cheers,
Miles

[gdalle]

However, one thing I was wondering if you would consider is designating a mapping from Symbol to the correct AbstractADType?

Indeed that used to be a widespread approach, but a Symbol is not enough to fully represent a choice of AD backend. That's why e.g. Optimization.jl switched to ADTypes.jl. The reasons are described in more detail in the ADTypes.jl README.
I fear that adding a layer of translation from Symbols to ADTypes.AbstractADType would only result in competing standards, as well as suboptimal default choices for all the things a Symbol cannot express.

Basically, I would prefer that my users do not need to install both Zygote.jl and DifferentiationInterface.jl (or ADTypes) in their environment to set Zygote as the backend.

Why is that? If they use Zygote.jl as the backend, it will be installed anyway. And AD backends are infinitely heavier than ADTypes.jl or DifferentiationInterface.jl. ADTypes.jl has zero dependency on Julia >= 1.9, and DifferentiationInterface.jl is very lightweight too. Besides, if you use them internally, they'll be in the Manifest.toml anyway.

which means the user has a more lightweight interface; they can just pass the symbol:

To me, the following interface is even better, because it does not introduce a separate Symbol mapping, and instead leverages ADTypes.jl directly. Keep in mind that users may already be familiar with this system if they have ever used a SciML package.

using ADTypes, SymbolicRegression, Zygote

options = Options(AutoZygote())

[gdalle]

And by the way, I'd be thrilled to help you get DifferentiationInterface.jl up and running with SymbolicRegression.jl. Don't hesitate to ping me for reviews or advice in that PR

[MilesCranmer]

Thanks!

It's not so much the extra dependency as SymbolicRegression.jl would of course already depend on DifferentiationInterface.jl if this goes through. The issue is I would be asking beginner users to install a third package in their top level environment, even after I have already asked them to install the AD backend package itself, just so they can set one of the parameters of the search.

Indeed that used to be a widespread approach, but a Symbol is not enough to fully represent a choice of AD backend.

I agree with this, in the example I gave above I allow both AbstractADType and a Symbol.

using DifferentiationInterface: lookup_backend, AbstractADType

function Options(autodiff_backend::Union{AbstractADType,Symbol})
    backend = autodiff_backend isa Symbol ? lookup_backend(autodiff_backend) : autodiff_backend

    return Options{typeof(backend)}(backend)
end

So advanced users can install ADTypes and customize the backend as they desire, but beginner users (which SymbolicRegression.jl gets a lot of from the Python frontend PySR) still can just pass in a :Zygote – which is probably good enough in 95% of cases.

(Even package extensions themselves I find to be a distraction for beginners, which is why PySR just automatically installs stuff as the user needs it – https://github.com/MilesCranmer/PySR/blob/89e991d9d7a5e34025bdbbe9076a8fd0892f04cc/pysr/julia_extensions.py#L8-L36 – rather than asking the user to install each package themselves. I just find it to be another barrier for people using my package.)

[gdalle]

I understand your point, and I do think it makes a lot of sense to offer an easier setting, especially for people coming from Python. But I don't think this Symbol mapping should be part of DifferentiationInterface.jl, if only because I try to make as few assumptions as possible.

The best of both worlds might be to automatically turn a Symbol into an ADTypes.AbstractADType as follows: :Package -> ADTypes.AutoPackage(). In most cases this will work (except for a few minor backends which always require arguments) and yield a reasonable backend object.
That way you don't have to maintain an explicit mapping, and you automatically benefit from e.g. new backends that might be added to ADTypes.jl.

julia> using ADTypes

julia> s = :Zygote
:Zygote

julia> auto_s = Symbol(:Auto, s)
:AutoZygote

julia> backend = @eval $auto_s()
AutoZygote()

julia> backend isa AbstractADType
true

There might be some subtleties as to which eval you should use / in which scope (I'm not very confident in metaprogramming) but you get the gist.

[MilesCranmer]

The reason I thought it might make sense to offer such a mechanism in this package is so that we can avoid having different ways of solving this. I am guessing that other packages may want to offer similar functionality in the future. So it would be nice to have a unified mapping from symbol to backend type.

Even something like

function Auto(key::Symbol; options...)
    backend_type = lookup_type(key)
    return backend_type(; options...)
end

Which would let me pass extra backend configuration options from the user.

Then I can import this function internally in my package, and allow the user to pass a symbol and autodiff configuration to my options constructor, without needing the user to manually install ADTypes.

[gdalle]

I don't think a unified mapping from Symbol to AbstractADType always makes sense. For all I know, other people may want a different mapping, or they may use symbols that don't correspond to an Auto struct at all, like :best_backend_for_linear_regression.
For DifferentiationInterface.jl, the ground truth is ADTypes.jl, and I want to avoid opinionated decisions on how to interpret a Symbol. If there are more people that need this mapping I might reconsider (maybe @Vaibhavdixit02 or @avikpal have opinions on this), but as a first step, I think parsing the Symbol as the second part of AutoPackage and adding kwargs is a reasonable thing to do on your end.
@adrhill thoughts?

[MilesCranmer]

I guess I feel like (symbol; kws...) -> Auto{symbol}(; kws...) is quite an innocent choice?

You could even make it like

function Auto{key}(; options...)
    backend_type = lookup_type(key)
    return backend_type(; options...)
end

rather than needing AutoEnzyme, AutoZygote, etc., which need to be explicitly imported.

So, to summarise my thoughts:

The issue with having the user install ADTypes.jl themselves is that it presents a small barrier to usage and more mental overhead. While it wouldn't slow down performance, it would mean another line in the user's Project.toml file, the extra effort of needing to run add ADTypes, as well as simply the user needing to figure this out in the first place, rather than simply passing in :Zygote which I can document pretty easily.

The two workarounds are:

1. Re-export ADTypes myself.
   • The issue with this is that I would either need to re-export all types (which is dangerous, in case of any future conflicts), or re-export the current set of types (meaning I would need to keep this up-to-date manually).
2. Write my own mapping from symbol => type.
   • This has the same issue of needing to keep this up-to-date manually with all of the types in ADTypes
   • This also introduces the problem of proliferating multiple standards for doing this, leading to a similar problem to the one this package was designed to solve.
   • If I get around this by using the ADType.eval(symbol("Auto", ad_backend)) trick, it is a bit dangerous (eval in general runs into https://en.wikipedia.org/wiki/Arbitrary_code_execution), and also might lead to issues with precompilation due to manual calls of eval

[adrhill]

The reason I thought it might make sense to offer such a mechanism in this package is so that we can avoid having different ways of solving this. I am guessing that other packages may want to offer similar functionality in the future.

From my POV, this is why ADTypes exists in the first place–it is the standard you are asking for. ADTypes are basically symbols with optional configuration options:

A natural approach is to use a keyword argument with e.g. Bool or Symbol values.
Let's see a few examples to understand why this is not enough:

• autodiff = true: ambiguous, we don't know which AD package should be used
• autodiff = :forward: ambiguous, there are several AD packages implementing both forward and reverse mode (and there are other modes beyond that)
• autodiff = :Enzyme: ambiguous, some AD packages can work both in forward and reverse mode
• autodiff = (:Enzyme, :forward): not too bad, but many AD packages require additional configuration (number of chunks, tape compilation, etc.)

A more involved struct is thus required, with package-specific parameters.
If every AD user develops their own version of said struct, it will ruin interoperability.
This is why ADTypes.jl provides a single set of shared types for this task, as an extremely lightweight dependency.
They are types and not enums because we need AD choice information statically to use it for dispatch.

As Guillaume mentioned, there are some very large packages putting their weight behind ADTypes as a unified interface to select backends:

Keep in mind that users may already be familiar with this system if they have ever used a SciML package.

Maybe reexporting ADTypes is an option if you don't want users to have to load the package explicitly? (EDIT: just saw you addressed this simultaneously.)

[MilesCranmer]

@adrhill oops, our messages crossed. Regarding this:

Maybe reexporting ADTypes is an option if you don't want users to have to load the package explicitly?

It's a bit tricky, because either I:

1. re-export everything in the package (which is a bit dangerous for a few reasons – including symbol conflicts, and the "spooky action at a distance" effect of not being explicit about what symbol comes from where), or
2. I keep a list of currently exported AD types myself, and regularly update it. (But this a bit of a pain)

[gdalle]

My issue is that if I make this translation available in DI (or in ADTypes itself for that matter), I somehow endorse it and encourage people to use it. But I don't want to, because there should only be one standard. Additionally, starting from a Symbol makes everything else type-unstable, unlike a properly defined AD backend.
Thus keeping a single standard seems worth the small mental overhead to me.

[adrhill]

Additionally, if you ever want to use our second order or sparse AD functionality¹, the amount of symbols will grow combinatorially.

Footnotes

1. Would be interesting to discuss whether this could have some applications in symbolic regression. We're still on the look-out for possible applications! ↩

[MilesCranmer]

Additionally, starting from a Symbol makes everything else type-unstable, unlike a properly defined AD backend.

Actually this wouldn't be an issue.

First, this wouldn't redefine anything, it's just making an mapping that packages can use for easy user-side APIs. But a user would always be able to pass in a full ADType instead to the options if they wish.

Second, if the symbol is fixed in some code, Julia would just inline it and it would be stable.

Additionally, if you ever want to use our sparse AD functionality1, the amount of symbols will grow combinatorially.

This also wouldn't be an issue – the symbol here would only be the type specification. Like this:

function Auto(symbol, options...; kws...)

    if symbol == :Zygote
        AutoZygote(options...; kws...)
    elseif symbol == :Enzyme
        AutoEnzyme(options...; kws...)
    end
end

So, in your example, it would be

Auto(
    :Sparse,
    Auto(:ForwardDiff);
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(),
)

which could be done without needing to re-export all the ADTypes.

[MilesCranmer]

By the way, just curious, why not have

const AutoZygote = Auto{:Zygote}
const AutoSparse = Auto{:Sparse}
const AutoEnzyme = Auto{:Enzyme}

instead of AutoZygote, AutoSparse, etc? This would also solve the problem as I can simply map the user's passed symbol to Auto{symbol} inside my constructor, rather than need the user to install and import AutoZygote themselves. They can of course pass options here as well, like

Auto{:Sparse}(
    Auto{:ForwardDiff}();
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(),
)

so you don't lose extensibility.

I mean in general I should say I don't feel too strongly about this. It just seems like it would be pretty useful for designing user-facing APIs.

I note that will have to do some symbol => type mapping anyways, because of PySR (users can't specify Julia types there). I felt like it would be bad practice to come up with my own AD symbol => AD type mapping when others may want to use one as well (not to mention there might be more AD types in the future, so should probably be centralized and kept up-to-date).