Any Chance to re-support Tracker.jl?

[patrickm663]

Hi all

Could support for something like Tracker.jl be re-introduced, like at a use-at-your-own-risk? If not, is there then a way to manually add in an AD library that isn't explicitly supported ie like Enzyme.jl? Reason being, I have BNNs I've been working on and Tracker.jl was working well enough, and was the fasted backend for the size of the network and data. Updating to 0.35 obviously broke some code, but I couldn't get Mooncake.jl to work as a drop-in replacement sampler.

I then tried replicating the Bayesian Neural Network tutorial directly in a clean environment but couldn't get Mooncake to work either. I'm away from my laptop at the moment, but I recall I would get a Stackoverflow error when running the sampler. Other AD libraries worked without modification to the code (although Zygote.jl required that I remove the f32 conversion). I haven't re-tested this, so it may be fine now :) I also tried loading AutoTracker() directly but couldn't get the sampler to work.

Sticking with the tutorial, I downgraded and tested Tracker.jl on the model. It benchmarked slightly behind ForwardDiff.jl (the NN only has ~20 parameters), but beat out Zygote.jl and ReverseDiff.jl. As soon as I upped the size of the NN (200+ parameters), Tracker.jl was a clear winner.

I do get Tracker.jl's not being actively maintained so I agree there shouldn't be tutorials on it, but it's not outright broken yet from my experience.

I would like to benchmark Mooncake.jl soon and see how it compares. I'd also be interested in seeing how performance is on the tutorial docs, since the output is currently suppressed.

Happy to share a MWE and the stacktrace I get when I'm back at my desk :)

[torfjelde]

Mooncake.jl is currently quite slow on Turing.jl due to a bug, so that performance should hopefully be much better very soon:) Ref: #2418

Regarding Tracker.jl, you can always post the MWE + stacktrace, but I do think it's unlikely that it'll be fixed unless it's obvious what's wrong + a very simple fix 😕

[patrickm663]

Thanks for the quick reply! Tracker's issue is I think just due to AutoTracker() getting dropped in the latest release from what I can tell, since it works in v0.34 then breaks in v0.35.

I'll have another look at Mooncake.jl to see if I can get it working for my code, then maybe open a separate Issue if it still breaks. Thanks for the note re the speed, I'll keep an eye on updates to the package :)

[willtebbutt]

If you want to attempt to use Tracker, you can always import ADTypes and use ADTypes.AutoTracker directly -- it should work :)

[willtebbutt]

Also, re stack overflow with Mooncake -- there's a known issue that I'm actively working on at the minute. It should be resolved tomorrow. Apologies for it not working out-of-the-box!

[patrickm663]

Thanks very much for the replies! I have been sticking to Turing v0.34.1 for the time being but will revist getting my set-up to work on the latest version of Turing now that I have some time over the holidays :) In my experience, loading ADTypes let me access AutoTracker() but I would get a long stacktrace of copyto!!() errors iirc when I then ran the sampler.

[willtebbutt]

Ooo that's annoying. I suspect that this is something to do with DifferentiationInterface (I believe that function is defined there), but you'd have to take another look at your stack trace and confirm.