Externalize gradient computations to DifferentiationInterface.jl?

[gdalle]

Hey there Avik!

As you may know, I have been busy developing DifferentiationInterface.jl, and it's really starting to take shape.
I was wondering if it would be useful for Lux.jl as a dependency, in order to support a wider variety of autodiff backends defined by ADTypes.jl?

Looking at the code, it seems the main spot where AD comes up (beyond the docs and tutorials) is Lux.Training:

Lux.jl/src/contrib/training.jl

Lines 93 to 106 in c27b9f5

	function compute_gradients(ad::ADTypes.AbstractADType, ::Function, _, ::TrainState)
	return __maybe_implemented_compute_gradients(ad)
	end
	function __maybe_implemented_compute_gradients(::T) where {T <: ADTypes.AbstractADType}
	throw(ArgumentError("Support for AD backend $(nameof(T)) has not been implemented yet!!!"))
	end
	function __maybe_implemented_compute_gradients(::ADTypes.AutoZygote)
	throw(ArgumentError("Load `Zygote` with `using Zygote`/`import Zygote` before using this function!"))
	end
	function __maybe_implemented_compute_gradients(::ADTypes.AutoTracker)
	throw(ArgumentError("Load `Tracker` with `using Tracker`/`import Tracker` before using this function!"))
	end

Gradients are only implemented in the extensions for Zygote and Tracker:

Lux.jl/ext/LuxZygoteExt.jl

Lines 7 to 14 in c27b9f5

	function Lux.Experimental.compute_gradients(
	::AutoZygote, objective_function::Function, data, ts::Lux.Experimental.TrainState)
	(loss, st, stats), back = Zygote.pullback(
	ps -> objective_function(ts.model, ps, ts.states, data), ts.parameters)
	grads = back((one(loss), nothing, nothing))[1]
	@set! ts.states = st
	return grads, loss, stats, ts
	end

Lux.jl/ext/LuxTrackerExt.jl

Lines 25 to 33 in c27b9f5

	function Lux.Experimental.compute_gradients(
	::AutoTracker, objective_function::Function, data, ts::Lux.Experimental.TrainState)
	ps_tracked = fmap(param, ts.parameters)
	loss, st, stats = objective_function(ts.model, ps_tracked, ts.states, data)
	back!(loss)
	@set! ts.states = st
	grads = fmap(Tracker.grad, ps_tracked)
	return grads, loss, stats, ts
	end

While DifferentiationInterface.jl is not yet ready or registered, it has a few niceties like Enzyme support which might pique your interest. I'm happy to discuss with you and see what other features you might need.

The main one I anticipate is compatibility with ComponentArrays.jl (JuliaDiff/DifferentiationInterface.jl#54), and I'll try to add it soon.

cc @adrhill

[avik-pal]

Yeah, I am all for purging that code, depending on DifferentiationInterface. The major things that need to work:

1. Tracker
2. ComponentArrays (probably just testing)

That should be enough to get a start.

[gdalle]

Tracker works and is tested, ComponentArrays is my next target

[avik-pal]

That's great. This week is a bit busy but I can try possibly early next month

[gdalle]

No rush! And always happy to help debug.
We're also chatting with Flux to figure out how best to support their use case, which is slightly more complex for lack of a dedicated parameter type like ComponentVector.

[avik-pal]

Do you require a vector input mandatorily? ComponentArrays has an overhead for smallish arrays (see #49), so having an fmap based API might be good, though that is not really a big blocker.

[gdalle]

At the moment yes. We're thinking about how to be more flexible in order to accommodate Flux's needs, you can track JuliaDiff/DifferentiationInterface.jl#87 to see how it evolves

[gdalle]

DI v0.3 should be out sometime next week (I've been busy with sparse Jacobians & Hessians), but I don't think I'll have much time in the near future to revamp the Lux tests. Still, I think it would make sense to offer DI at least as a high level interface, even if it is not yet used in the package internals / tests. It might also help you figure out #605

[gdalle]

Note that for DI to work in full generality with ComponentArrays, I need SciML/ComponentArrays.jl#254 to be fixed. Otherwise Jacobians and Hessians will stay broken (the rest, in particular gradient, is independent from stacking)

[avik-pal]

Yes I want to roll it out first as a high level interface when the inputs are AbstractArray rather than for testing. I still want to support arbitrary structured parameters as input for now.