add Tapir support

[Red-Portal]

No description provided.

[yebai]

@sunxd3 can you help complete this integration?

[sunxd3]

After I finish the work on hands, will work on this!

[yebai]

I formatted the code on the master branch and this PR so the code changes are more readable. Unfortunately, this means some manual merge efforts are required for other PRs, like #67.

[yebai]

@willtebbutt It seems the only remaining issue is the lack of a rule for Intrinsics.fptrunc. Can this be quickly added?

https://github.com/TuringLang/AdvancedVI.jl/actions/runs/10321643466/job/28575032583?pr=71#step:6:1408

[willtebbutt]

Just had a look at the source -- it looks like this is happening in the rule that Tapir.jl is deriving for randn!. Shoudn't be a problem to add a rule because this intrinsic seems fairly safe (it's converting a Float64 into a Float32).

[Red-Portal]

I formatted the code on the master branch and this PR so the code changes are more readable. Unfortunately, this means some manual merge efforts are required for other PRs, like #67.

@yebai Thanks! I'll deal with 67 myself.

[Red-Portal]

@willtebbutt BTW, in this PR, I had to modify the package so that we carry around the rrule. I personally don't want to do this, is there a Tapir native way? DifferentiationInterface unfortunately doesn't support auxiliary input yet, so I can't rely on it for this (unless we wait until it gets implemented.)

[willtebbutt]

@willtebbutt BTW, in this PR, I had to modify the package so that we carry around the rrule. I personally don't want to do this, is there a Tapir native way? DifferentiationInterface unfortunately doesn't support auxiliary input yet, so I can't rely on it for this (unless we wait until it gets implemented.)

Alas, there is not presently -- you'll have to continue to pass it around.

[willtebbutt]

@willtebbutt It seems the only remaining issue is the lack of a rule for Intrinsics.fptrunc. Can this be quickly added?

This is now resolved in v0.2.33 .

[willtebbutt]

Some of the remaining problems are to do with a different intrinsic, Core.Intrinsics.fpext. A version of Tapir.jl which addresses this should be available in an hour or so.

[willtebbutt]

I'm pretty sure that my comment will unblock this PR, but not everything will pass until v0.2.38 of Tapir.jl is available -- PR here

[yebai]

Suggested change

	DiffResults.gradient!(out, last(g))
	DiffResults.gradient!(out, g[2])

[yebai]

@willtebbutt, to clarify, we don't need this change. Is that correct?

[yebai]

All tests are passing now!

[yebai]

@willtebbutt, can you adapt the Bijectors setup so we don't need to comment out 1.7?

[yebai]

#67 introduces a few breaking changes to the value_and_gradient! interface, which is not yet incorporated in this PR.

[willtebbutt]

@Red-Portal is your plan to continue working on this PR, or should I take a look at finishing it off now that everything seems to be working on Tapir.jl's end?

[Red-Portal]

Hi all, sorry for the confusion caused by #67 . I was a little surprised that it broke some things because it wasn't changing the interface at all. Well turns out it is this PR that is breaking the interface (purely to carry around the rrules), but the master branch is still using the old interface, which is what is causing the conflicts.

Now that #67 is out of the way, I will be able to to spend more time on this PR. I'll take a look from now on. But I really want to do something about carrying around the rrules. Maybe I should look into Memoization?

[yebai]

Well turns out it is this PR that is breaking the interface (purely to carry around the rrules), but the master branch is still using the old interface, which is what is causing the conflicts.

It would be good to avoid heuristics like Memoization and pass cached rules explicitly. Can we upgrade everything to this PR's interface?

[Red-Portal]

If Memoization works for this, we can go back to the old interface and contain the memoization stuff in Tapir's extension. I personally think this is the way to go if Memoization works since passing around rrules is just way too complicated and only necessary because of Tapir.

[yebai]

I don’t have good experience with memorisation tricks; they lead to subtle bugs and are often hard to reason about. Maybe @willtebbutt have thought about this?

[willtebbutt]

I would also be wary of using Memoization here. At least one issue that could result in not-entirely-obvious cache invalidation is method redefinition. The global method table handles this well obviously, but I'm not sure it's entirely trivial to do yourself (hence the fact that Tapir.jl doesn't yet do it).

Re Tapir.jl being the only backend to require a preparation step: is it not the case that ReverseDiff.jl requires a cache when operating in compiled mode?

[Red-Portal]

Hi @willtebbutt !

At least one issue that could result in not-entirely-obvious cache invalidation is method redefinition.

Would this happen? We only use global functions now and the only things that change are auxiliary inputs. So if we are thinking of the same thing, I think it should be okay.

Re Tapir.jl being the only backend to require a preparation step: is it not the case that ReverseDiff.jl requires a cache when operating in compiled mode?

I don't think we'll be able to use compiled tapes because they are incompatible with changing auxiliary inputs.

[yebai]

There are a few details to consider carefully. I suggest staying away from memorization heuristics; the cons often outweigh the minor efforts that we could save (from using memorization) when it comes to finding a proper solution.

I hope that @willtebbutt can take over from here to ensure a robust AD interface for Tapir and other packages.

[willtebbutt]

Would this happen? We only use global functions now and the only things that change are auxiliary inputs. So if we are thinking of the same thing, I think it should be okay.

This can definitely happen -- Julia's world age system is used to track this, so if you were to do caching you would definitely need to take it into account.

[Red-Portal]

Hi @willtebbutt @yebai , okay after seeing the discussion, I guess carrying around rrules is the best we can do for now. Then I agree with this until DifferentiationInferface catches up or similar.

[yebai]

Closed in favour of #86