Re-add expansion of Enzyme.@import_{r,f}rule to BijectorsEnzymeExt

[penelopeysm]

Since Enzyme now works, but Julia hasn't been fixed, this re-enables the extension. See #339 (comment).

In the process, it aso re-enables tests that are no longer broken.

[devmotion]

IMO that's incomprehensible. What's the reason for not just implementing the forward and reverse-mode Enzyme rules? The function and the partial derivatives are very simple, and a direct implementation is likely faster and avoids unnecessary overhead. As also mentioned in the docstrings of @import_rrule and @import_frule, these macros most likely do not give you the most performant implementation (and they might lead to broken rules but I assume it's fine in this simple case here).

[penelopeysm]

I agree it is very incomprehensible, but we discussed this previously (on Slack) and none of us really have the time to figure out how to rewrite the rule ourselves. (I definitely don't, I don't know if anyone else's schedule has freed up, though I doubt it.) @wsmoses suggested that we could, in the worst case, run macroexpand, so that's what I've done, and it's only really intended as a stopgap measure until 1.11.2 is released, at which point we can get rid of all of this.

[devmotion]

it's only really intended as a stopgap measure until 1.11.2 is released, at which point we can get rid of all of this.

My point is that even when the Pkg issues are fixed, I think ideally you would use neither @import_frule nor @import_rrule.

[penelopeysm]

I will pass that on :)

[penelopeysm]

But maybe @willtebbutt or @mhauru one of you are better placed to do this (at some point in time)?

[yebai]

Regardless of how we want to implement these rules, Bijectors are only users of the rules defined here.

So, it feels more natural for these rules to be hosted as an EnzymeRootsExt extension, living in either the Enzyme.jl or the Roots.jl repo.

[yebai]

@devmotion, since you wrote the ChainRulesCore rule for Roots.jl, would you be happy to do the same for Enzyme, or coach @mhauru to do it?

[yebai]

Here are the rules for ReverseDiff, which can be adapted for Enzyme:

Bijectors.jl/ext/BijectorsReverseDiffExt.jl

Lines 246 to 258 in 2ee8a5d

	function find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:TrackedReal}
	return track(find_alpha, wt_y, wt_u_hat, b)
	end
	@grad function find_alpha(wt_y::TrackedReal, wt_u_hat::TrackedReal, b::TrackedReal)
	α = find_alpha(value(wt_y), value(wt_u_hat), value(b))
	∂wt_y = inv(1 + wt_u_hat * sech(α + b)^2)
	∂wt_u_hat = -tanh(α + b) * ∂wt_y
	∂b = ∂wt_y - 1
	find_alpha_pullback(Δ::Real) = (Δ * ∂wt_y, Δ * ∂wt_u_hat, Δ * ∂b)
	return α, find_alpha_pullback
	end

[penelopeysm]

Is it sensible to merge this first and keep track of the proper rule-writing in another issue?

Right now we have a missing rule and that will make things break.

If we'd rather let things sit broken for now, that's fine with me too! Just thought to prod.

[yebai]

It doesn't matter much since the code is related to normalising flows, which Turing does not use. It is more for users who directly work with Bijectors.

I'd strongly encourage that these rules got correctly fixed and moved to an EnzymeRoots extension instead of being kept here.

[devmotion]

@devmotion, since you wrote the ChainRulesCore rule for Roots.jl, would you be happy to do the same for Enzyme, or coach @mhauru to do it?

Well, one could do that. But that's again too generic for Bijectors IMO. The only interest here is one specific function, so only a rule for this specific function is needed. I assume generally performance of such a specific rule will be better than performance of generic Roots.jl rules.

I should have some time this evening, I'll open a draft PR and see how far I can get.