Found unhandled active variable in tuple splat

[aplavin]

Constructing tuples larger than a certain length throws this error that I couldn't even understand. MWE:

using Enzyme

function f(x)
	res = NTuple{40}(x)
	return res[1]
end

x = ones(40)
δδx = zeros(40)

(_, y) = Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(x, δδx))

# error:
AssertionError: Found unhandled active variable in tuple splat, jl_apply_iterate Tuple{Float64}
error_if_active_iter(::Base.RefValue{Tuple{Float64}})@jitrules.jl:775
_totuple@tuple.jl:410[inlined]
_totuple@tuple.jl:409[inlined]
Tuple@tuple.jl:391
...

The same code works if 40 is replaced with, say, 20.

How can this be fixed?

[wsmoses]

So that error message indicates that you're trying to differentiate a type unstable piece of Julia code, which we haven't added support for yet. In particular it's saying Enzyme does not yet have support for differentiating the type unstable tuple splat operator (aka (x...,) where x is an array). In particular this is the Julia runtime function jl_apply_iterate. We actually do have partial support for this, but only for duplicated variables (aka tuple containing vectors or ref etc, not tuples containing active variables like floats). This is on our plate to add support for, but in the interim the usual advice of making your code type stable applies (Enzyme presently only has partial support for type unstable code. Whether your code is type stable or not is unfortunately a Julia not Enzyme-level question, and there are some tools to identify this, like code_warntype). One way that you may be able to rewrite this to be type stable (on mobile so haven't tried), is something like res = (x[1], x[2]) etc for whatever the size of the tuple is. Obviously this may become unwield for larger tuples so you may want to auto generate this through the use of a generator or macro. Making your code type stable will not just enable Enzyme support, but also significantly improved performance of the code as well.

…

On Sun, Jan 28, 2024, 9:24 AM Alexander Plavin ***@***.***> wrote: Constructing tuples larger than a certain length throws this error that I couldn't even understand. MWE: using Enzyme function f(x) res = NTuple{40}(x) return res[1]end x = ones(40) δδx = zeros(40) (_, y) = Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(x, δδx)) # error: AssertionError: Found unhandled active variable in tuple splat, jl_apply_iterate ***@***.***:775 ***@***.***:410[inlined] ***@***.***:409[inlined] ***@***.***:391... The same code works if 40 is replaced with, say, 20. How can this be fixed? — Reply to this email directly, view it on GitHub <#1255>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXBSKTDEDJNGZBEI2QTYQZNTJAVCNFSM6AAAAABCOG42E2VHI2DSMVQWIX3LMV43ASLTON2WKOZSGEYDIMJYG4ZTCOI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[aplavin]

The example code should be fully type-stable though, it just creates a tuple with comp-time known length and eltype. Indeed, code_warntype doesn't show any instabilities.

[wsmoses]

I’ll take a deeper look when I get to a computer, but the error message can only come from trying to diffeeentiate the type unstable Julia splat operator (jl_apply_iterate), so that’s what the Julia compiler generated.

…

On Sun, Jan 28, 2024 at 11:09 AM Alexander Plavin ***@***.***> wrote: The example code should be fully type-stable though, it just creates a tuple with comp-time known length and eltype. Indeed, code_warntype doesn't show any instabilities. — Reply to this email directly, view it on GitHub <#1255>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXFGWIQ45RHTMI66WDLYQZZ5LAVCNFSM6AAAAABCOG42E2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMJTGY2DMNRRG4> . You are receiving this because you commented.Message ID: ***@***.***>

[wsmoses]

If helpful I think that by default warntype doesn’t recurse, and from the stacktrace you gave it looked like the type instability was from the inside of the tuple constructor.

…

On Sun, Jan 28, 2024 at 11:42 AM Billy Moses ***@***.***> wrote: I’ll take a deeper look when I get to a computer, but the error message can only come from trying to diffeeentiate the type unstable Julia splat operator (jl_apply_iterate), so that’s what the Julia compiler generated. On Sun, Jan 28, 2024 at 11:09 AM Alexander Plavin < ***@***.***> wrote: > The example code should be fully type-stable though, it just creates a > tuple with comp-time known length and eltype. Indeed, code_warntype doesn't > show any instabilities. > > — > Reply to this email directly, view it on GitHub > <#1255>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/AAJTUXFGWIQ45RHTMI66WDLYQZZ5LAVCNFSM6AAAAABCOG42E2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMJTGY2DMNRRG4> > . > You are receiving this because you commented.Message ID: > ***@***.***> >

[wsmoses]

If you’d like some additional debugging tips in the interim, If you set Enzyme.API.printall!(true) right after loading Enzyme it will print the input IR we get from Julia followed by the differentiated IR. If you see the jl_apply_iterate there that’s the type unstable call you’re hitting. code_llvm presumably would show it too, though it may only show the outermost function and have a call to a function with the type unstable call.

…

On Sun, Jan 28, 2024 at 11:44 AM Billy Moses ***@***.***> wrote: If helpful I think that by default warntype doesn’t recurse, and from the stacktrace you gave it looked like the type instability was from the inside of the tuple constructor. On Sun, Jan 28, 2024 at 11:42 AM Billy Moses ***@***.***> wrote: > I’ll take a deeper look when I get to a computer, but the error message > can only come from trying to diffeeentiate the type unstable Julia splat > operator (jl_apply_iterate), so that’s what the Julia compiler generated. > > On Sun, Jan 28, 2024 at 11:09 AM Alexander Plavin < > ***@***.***> wrote: > >> The example code should be fully type-stable though, it just creates a >> tuple with comp-time known length and eltype. Indeed, code_warntype doesn't >> show any instabilities. >> >> — >> Reply to this email directly, view it on GitHub >> <#1255>, >> or unsubscribe >> <https://github.com/notifications/unsubscribe-auth/AAJTUXFGWIQ45RHTMI66WDLYQZZ5LAVCNFSM6AAAAABCOG42E2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMJTGY2DMNRRG4> >> . >> You are receiving this because you commented.Message ID: >> ***@***.***> >> >

[wsmoses]

Oh also this is in the wrong repository. This is the core Enzyme code, not the Julia bindings.

I mean you're welcome to use Enzyme on C++, JaX, etc where this error cannot occur (being Julia-specific), but probably this isn't where this should go. Moving the issue.

[wsmoses]

Oh also even though codewarntype doesn't yell, it is clearly type unstable.

julia> @code_warntype f(x)
MethodInstance for f(::Vector{Float64})
  from f(x) @ Main REPL[4]:1
Arguments
  #self#::Core.Const(f)
  x::Vector{Float64}
Locals
  res::NTuple{40, Float64}
Body::Float64
1 ─ %1 = Core.apply_type(Main.NTuple, 40)::Core.Const(NTuple{40, T} where T)
│        (res = (%1)(x))
│   %3 = Base.getindex(res, 1)::Float64
└──      return %3

%1 = Core.apply_type(Main.NTuple, 40)::Core.Const(NTuple{40, T} where T) <- this returns an abstract type
as is (res = (%1)(x)) <- this is type unstable

If you change it to be slightly more type stable (and not an abstract type):

using Enzyme

function f(x)
   res = NTuple{40, Float64}(x)
   return res[1]
end

x = ones(40)
δδx = zeros(40)

(_, y) = Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(x, δδx))

julia> @code_warntype f(x)
MethodInstance for f(::Vector{Float64})
  from f(x) @ Main REPL[10]:1
Arguments
  #self#::Core.Const(f)
  x::Vector{Float64}
Locals
  res::**Tuple{Vararg{Float64}}**.  <- this is type unstable
Body::Float64
1 ─ %1 = Core.apply_type(Main.NTuple, 40, Main.Float64)::Core.Const(NTuple{40, Float64})
│        (res = (%1)(x))
│   %3 = Base.getindex(res, 1)::Float64
└──      return %3

[wsmoses]

If you actually make this type stable, like below, it does not error:

using Enzyme


function f(x)
	res = ntuple(Val(40)) do i
		Base.@_inline_meta
		x[i]
	end
	return res[1]
end

x = ones(40)
δδx = zeros(40)

(_, y) = Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(x, δδx))

((nothing,), 1.0)

[wsmoses]

Now closing as duplicate of #1235

If you'd like, feel free to open an issue on the Julia compiler itself (https://github.com/JuliaLang/julia) for the NTuple constructor not being type stable

[aplavin]

%1 = Core.apply_type(Main.NTuple, 40)::Core.Const(NTuple{40, T} where T) <- this returns an abstract type
as is (res = (%1)(x)) <- this is type unstable

code_warntype for my original example looks totally the same for 20 and for 40 (aside from the number of course). However, for N=20 Enzyme can autodiff through, for N=40 it cannot...

[wsmoses]

Yeah this is a question better posed to JuliaLang/julia at this point. All I can tell you is that the llvm the Julia compiler makes is type unstable by the time Enzyme gets it for your test case.

Feel free to cc me though.