Unroll _transform_tuple to fix Enzyme.autodiff on tuples of length ≥ 33

[jlperla]

Replace the Base.tail-recursive _transform_tuple with a @generated straight-line unroll — same outputs bit-for-bit, but the typed IR no longer contains a self-invoke, which is what Enzyme.autodiff (Forward and Reverse) trips on at tuple length ≥ 33 with AssertionError("conv == 37") (EnzymeAD/Enzyme.jl#3104).

[scheidan]

Any chance that this will be merged here? I understand that the real fix should be on Enzyme's side, but that may be much harder.

Thanks!

PS: The is my real world MWE that lead me finally to this PR; maybe it is useful for someone.

using Distributions
using Enzyme
using TransformVariables

N = 33
dists = ntuple(i -> LogNormal(0.0, 1.0), N)
dists = NamedTuple{ntuple(i -> Symbol("x", i), N)}(dists)

function prior_transform(priors)
    transforms = map(priors) do prior
        left, right = extrema(support(prior))
        left = isinf(left) ? -TransformVariables.∞ : left
        right = isinf(right) ? TransformVariables.∞ : right
        TransformVariables.as(Real, left, right)
    end
    TransformVariables.as(transforms)
end

trans = prior_transform(dists)
q = fill(-0.1, TransformVariables.dimension(trans))

foo(q) = sum(values(TransformVariables.transform(trans, q)))

Enzyme.gradient(Enzyme.Reverse, foo, q) # AssertionError: conv == 37

[tpapp]

@jlperla, thanks for this, @scheidan, thanks for the ping. I apologize for the delay in reviewing this.

It is not strictly equivalent as, AFAIK, built-ins do not necessarily unroll above a certain tuple length. But given that the intention of using a tuple is to get type-stable code, I don't see a problem with this here. Also, EnzymeAD/Enzyme.jl#3104 indicates that this is an issue on the Julia side, so fixing it on our end may be the best option for now.

@devmotion, this is fine with me, do you have any comments?

[tpapp]

(closing and reopening to make CI run)

[devmotion]

Why is a separate idx variable needed? Couldn't we just operate with index?

[tpapp]

LGTM, thanks!

[devmotion]

same outputs bit-for-bit

Just wanted to mention: This is wrong, I actually saw downstream test failures due to this PR. Previously, summation was performed using (basically) foldr whereas this PR uses foldl, which does of course generally not give exactly the same results.

[jlperla]

I was waiting for CI to run. Did it?

[devmotion]

In this PR? Sure, it ran before it was merged. My point was merely that it broke my downstream CI due to not yielding "same outputs bit-for-bit", and I wanted to point out for any future reader of this PR that this was an incorrect claim in the initial comment above.

[tpapp]

same outputs bit-for-bit

I just want to clarify that this is not something this package ever promised. AFAIK very few packages in the Julia ecosystem have that kind of commitment. @devmotion, thanks for pointing this out though.

downstream test failures due to this PR

I am sorry to hear this, but if they were comparing exact output, those were the wrong kind of tests.

[devmotion]

No, they were not comparing to exact TransformVariables output. The test failure was caused by different likelihood values of MLE estimates, apparently the tiny difference in the transform was sufficient to cause slightly different optimization trajectories. Not a big problem, of course, but that made me realize the incorrect claim above.