AutoEnzyme tries to differentiate Boolean Mask

[mmikhasenko]

The Enzyme was the last of AD that gave up on my problems among all ADs.
A bug is subtle, it seems only reproduced by combining two PartitionMasks

ERROR: InexactError: Bool(1.559401f0)
Stacktrace:
  [1] Bool
    @ ./float.jl:251 [inlined]
  [2] convert
    @ ./number.jl:7 [inlined]
  [3] _setindex_scalar!
    @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:3183
  [4] setindex!
    @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:3180 [inlined]
  [5] reverse
    @ ~/.julia/packages/Enzyme/9dy6G/src/internal_rules.jl:854 [inlined]
  [6] mul!
    @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:253 [inlined]
  [7] *
    @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:60 [inlined]
  [8] combine
    @ ~/Documents/ML.CAT/Bijectors.jl/src/bijectors/coupling.jl:125 [inlined]

pointing to the combine method in Bijectors.jl.

I did my best to simplify, the example that reproduce the error.

using Bijectors
using NormalizingFlows.Optimisers
using DifferentiationInterface
import Flux: Dense
import Enzyme

struct Trivial{T} <: Bijectors.Bijector
    params::AbstractVector{T}
end

Bijectors.transform(b::Trivial, x) = with_logabsdet_jacobian(b, x)[1]
Bijectors.logabsdetjac(b::Trivial, x) = with_logabsdet_jacobian(b, x)[2]
Bijectors.with_logabsdet_jacobian(b::Trivial, x::Real) = x, b.params[1]
Bijectors.with_logabsdet_jacobian(b::Trivial, x::AbstractVector) = x, b.params[1]
Bijectors.with_logabsdet_jacobian(ib::Inverse{<:Trivial}, y::Real) = y, -ib.orig.params[1]
Bijectors.with_logabsdet_jacobian(ib::Inverse{<:Trivial}, y::AbstractArray) = y, -ib.orig.params[1]

flow = let
    q₀ = Distributions.Product([Uniform(0, 1), Uniform(0, 1)])
    mask12 = Bijectors.PartitionMask(2, [1], [2])
    mask21 = Bijectors.PartitionMask(2, [2], [1])
    layers = [
        Coupling(Trivial ∘ Dense(1, 2), mask12),
        Coupling(Trivial ∘ Dense(1, 2), mask21),
    ]
    ts = reduce(∘, layers)
    transformed(q₀, ts)
end

let
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -pdf(re(θ), [0.5, 0.5])
    loss(θ)
end # works

let # fails
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -pdf(re(θ), [0.5, 0.5])

    ADbackend = AutoEnzyme(;
        mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
        function_annotation = Enzyme.Const)

    prep = prepare_gradient(loss, ADbackend, θ)
    DifferentiationInterface.gradient(loss, prep, ADbackend, θ)
end

An observation: the code works with 1 layer, and breaks with two layers

[gdalle]

Do all other autodiff backends work?

[mmikhasenko]

none of the others, for various reasons,
e.g. Mooncake: chalk-lab/Mooncake.jl#660

ForwardDiff and ReverseDiff gave up rather early because of the special types Dual and TrackedReal do not play with Flux

[mmikhasenko]

Btw: I tried to reimplement PartitionMask without using SparseArrays, and multiplication.
It overcomes this particular problem.

Code that works for three years (42 parameters)

Note: RationalQuadraticSpline are from mmikhasenko/Bijectors.jl#1

using Bijectors
using NormalizingFlows.Optimisers
using DifferentiationInterface

import Flux: Dense
using Enzyme: Enzyme

using Random
rng = MersenneTwister(42)

struct TwoDMask <: Bijectors.AbstractMask
    i::Int
    j::Int
end

Bijectors.combine(m::TwoDMask, x_1, x_2, x_3) = m.i < m.j ? vcat(x_1, x_2) : vcat(x_2, x_1)
Bijectors.partition(m::TwoDMask, x) = (x[m.i:m.i], x[m.j:m.j], x[m.j:m.j])

flow = let
    q₀ = Distributions.Product([Uniform(0, 1), Uniform(0, 1)])
    mask12 = TwoDMask(1, 2)
    mask21 = TwoDMask(2, 1)
    layers = [
        Coupling(RationalQuadraticSpline ∘ Dense(1, 7), mask12),
        Coupling(RationalQuadraticSpline ∘ Dense(1, 7), mask21),
        Coupling(RationalQuadraticSpline ∘ Dense(1, 7), mask12),
    ]
    ts = reduce(∘, layers)
    transformed(q₀, ts)
end

let
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -pdf(re(θ), Float32[0.5, 0.5])
    loss(θ)
end # works

let # works with 3 layers
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -pdf(re(θ), Float32[0.5, 0.5])

    ADbackend = AutoEnzyme(;
        mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
        function_annotation = Enzyme.Const)

    prep = prepare_gradient(loss, ADbackend, θ)
    DifferentiationInterface.gradient(loss, prep, ADbackend, θ)
end

When I put not three, but four layers, the AD falls short with the StackOverFlow.

[gdalle]

Can you reproduce this without DI, using Enzyme.gradient?

[mmikhasenko]

I tried,

let # fails
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -logpdf(re(θ), rand(Float32, 2))
    Enzyme.gradient(Enzyme.Reverse, loss, θ)
end

how do I go around this error?

ERROR: 
If you are using Enzyme by selecting the `AutoEnzyme` object from ADTypes, you may want to try setting the `function_annotation` option as follows:

        AutoEnzyme(; function_annotation=Enzyme.Duplicated)

This hint appears because DifferentiationInterface and Enzyme are both loaded. It does not necessarily imply that Enzyme is being called through DifferentiationInterface.

Function argument passed to autodiff cannot be proven readonly.
If the the function argument cannot contain derivative data, instead call autodiff(Mode, Const(f), ...)
See https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Activity-of-temporary-storage for more information.
The potentially writing call is   store {} addrspace(10)* %.fca.0.0.0.0.extract, {} addrspace(10)* addrspace(10)* %.innerparm.sroa.0.0..sroa_cast, align 8, !dbg !92, !alias.scope !102, !noalias !103, using   %.fca.0.0.0.0.extract = extractvalue [1 x { { [1 x {} addrspace(10)*], [2 x { { {} addrspace(10)*, { {} addrspace(10)*, {} addrspace(10)* } }, [3 x { i64, i64, {} addrspace(10)*, {} addrspace(10)*, {} addrspace(10)* }] }] }, { [1 x {} addrspace(10)*], [2 x { { { i64, i64 } } }] }, i64 }] %0, 0, 0, 0, 0, !dbg !12, !enzyme_type !73, !enzymejl_byref_MUT_REF !11, !enzymejl_source_type_Vector\7BUniform\7BFloat64\7D\7D !11

Stacktrace:

I check I do not expose DifferentiationInterface

import Bijectors: Coupling, RationalQuadraticSpline, PartitionMask, transformed
using Bijectors.Distributions
import NormalizingFlows.Optimisers

import Flux: Dense
using Enzyme

[gdalle]

Have you tried wrapping f into Const(f) as the error message suggests?
It is the native Enzyme equivalent of setting function_annotation=Enzyme.Const inside AutoEnzyme

[mmikhasenko]

Ah, overlooked it,

let # fails
    θ, re = Optimisers.destructure(flow)
    loss(θ) = -logpdf(re(θ), rand(Float32, 2))
    Enzyme.gradient(Enzyme.Reverse, Const(loss), θ)
end

ERROR: Constant memory is stored (or returned) to a differentiable variable.
As a result, Enzyme cannot provably ensure correctness and throws this error.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Runtime-Activity).
If Enzyme should be able to prove this use non-differentable, open an issue!
To work around this issue, either:
 a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
 b) set the Enzyme mode to turn on runtime activity (e.g. autodiff(set_runtime_activity(Reverse), ...) ). This will maintain correctness, but may slightly reduce performance.
Mismatched activity for:   store {} addrspace(10)* %.fca.0.0.0.0.extract, {} addrspace(10)* addrspace(10)* %.innerparm.sroa.0.0..sroa_cast, align 8, !dbg !97, !alias.scope !107, !noalias !108 const val:   %.fca.0.0.0.0.extract = extractvalue [1 x { { [1 x {} addrspace(10)*], [2 x { { {} addrspace(10)*, { {} addrspace(10)*, {} addrspace(10)* } }, [3 x { i64, i64, {} addrspace(10)*, {} addrspace(10)*, {} addrspace(10)* }] }] }, { [1 x {} addrspace(10)*], [2 x { { { i64, i64 } } }] }, i64 }] %0, 0, 0, 0, 0, !dbg !17, !enzyme_type !78, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_Vector\7BUniform\7BFloat64\7D\7D !0
 value=Unknown object of type Vector{Uniform{Float64}}
 llvalue=  %.fca.0.0.0.0.extract = extractvalue [1 x { { [1 x {} addrspace(10)*], [2 x { { {} addrspace(10)*, { {} addrspace(10)*, {} addrspace(10)* } }, [3 x { i64, i64, {} addrspace(10)*, {} addrspace(10)*, {} addrspace(10)* }] }] }, { [1 x {} addrspace(10)*], [2 x { { { i64, i64 } } }] }, i64 }] %0, 0, 0, 0, 0, !dbg !17, !enzyme_type !78, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_Vector\7BUniform\7BFloat64\7D\7D !0

Stacktrace:
  [1] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126
  [2] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177
  [3] execute
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:55
  [4] #fmap#41
    @ ~/.julia/packages/Functors/LbNAu/src/maps.jl:11
  [5] fmap
    @ ~/.julia/packages/Functors/LbNAu/src/maps.jl:3
  [6] #_rebuild#58
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:96
  [7] _rebuild
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:94
  [8] Restructure
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:59
  [9] loss
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:47
 [10] loss
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:0

Stacktrace:
  [1] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
  [2] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
  [3] execute
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:55 [inlined]
  [4] #fmap#41
    @ ~/.julia/packages/Functors/LbNAu/src/maps.jl:11 [inlined]
  [5] fmap
    @ ~/.julia/packages/Functors/LbNAu/src/maps.jl:3 [inlined]
  [6] #_rebuild#58
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:96 [inlined]
  [7] _rebuild
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:94 [inlined]
  [8] Restructure
    @ ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:59 [inlined]
  [9] loss
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:47 [inlined]
 [10] loss
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:0 [inlined]
 [11] augmented_julia_loss_32927_inner_1wrap
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:0
 [12] macro expansion
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5610 [inlined]
 [13] enzyme_call
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5144 [inlined]
 [14] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5083 [inlined]
 [15] autodiff
    @ ~/.julia/packages/Enzyme/9dy6G/src/Enzyme.jl:408 [inlined]
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/9dy6G/src/sugar.jl:324 [inlined]
 [17] gradient(::ReverseMode{false, false, false, FFIABI, false, false}, ::Const{var"#loss#13"{…}}, ::Vector{Float32})
    @ Enzyme ~/.julia/packages/Enzyme/9dy6G/src/sugar.jl:262
 [18] top-level scope
    @ ~/Documents/ML.CAT/Bijectors.jl/test/struggle.jl:48
Some type information was truncated. Use `show(err)` to see complete types.

[gdalle]

Can you also follow the indication in the error message to turn on runtime activity? After all, you did it for the DI version 😉
I'm trying to see if DI actually plays a role here or not. If not, this issue would be more at home on the Enzyme repo.

[gdalle]

So here's a fully DI-free MWE. I don't get the same error as you, maybe it's because I'm on different package versions, or maybe it's because I made your code type-stable with a few tweaks (making global variables constant, fixing type parameter of your struct). If that's okay with you I'm gonna close the issue here, and invite you to reopen it on the Enzyme side?

using Bijectors
using NormalizingFlows.Optimisers
import Flux: Dense
using Enzyme: Enzyme
Enzyme.Compiler.VERBOSE_ERRORS[] = true

struct Trivial{T,V<:AbstractVector{T}} <: Bijectors.Bijector
    params::V
end

Bijectors.transform(b::Trivial, x) = with_logabsdet_jacobian(b, x)[1]
Bijectors.logabsdetjac(b::Trivial, x) = with_logabsdet_jacobian(b, x)[2]
Bijectors.with_logabsdet_jacobian(b::Trivial, x::Real) = x, b.params[1]
Bijectors.with_logabsdet_jacobian(b::Trivial, x::AbstractVector) = x, b.params[1]
Bijectors.with_logabsdet_jacobian(ib::Inverse{<:Trivial}, y::Real) = y, -ib.orig.params[1]

function Bijectors.with_logabsdet_jacobian(ib::Inverse{<:Trivial}, y::AbstractArray)
    return y, -ib.orig.params[1]
end

const flow = let
    q₀ = Distributions.Product([Uniform(0, 1), Uniform(0, 1)])
    mask12 = Bijectors.PartitionMask(2, [1], [2])
    mask21 = Bijectors.PartitionMask(2, [2], [1])
    layers = [
        Coupling(Trivial ∘ Dense(1, 2), mask12), Coupling(Trivial ∘ Dense(1, 2), mask21)
    ]
    ts = reduce(∘, layers)
    transformed(q₀, ts)
end

const θ, re = Optimisers.destructure(flow)

loss(θ) = -pdf(re(θ), [0.5, 0.5])

Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(loss), θ)

Details

julia> Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(loss), θ)
ERROR: Enzyme compilation failed due to an internal error.
 Please open an issue with the code to reproduce and full error log on github.com/EnzymeAD/Enzyme.jl
 To toggle more information for debugging (needed for bug reports), set Enzyme.Compiler.VERBOSE_ERRORS[] = true (default false)
Current scope: 
; Function Attrs: mustprogress willreturn
...


Stacktrace:
 [1] -
   @ ./float.jl:489
 [2] with_logabsdet_jacobian
   @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:18
 [3] with_logabsdet_jacobian
   @ ~/.julia/packages/Bijectors/XFXMR/src/bijectors/coupling.jl:226

Stacktrace:
  [1] julia_error(msg::String, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/errors.jl:461
  [2] julia_error(cstr::Cstring, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/errors.jl:299
  [3] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, forceAnonymousTape::Bool, runtimeActivity::Bool, strongZero::Bool, width::Int64, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/9dy6G/src/api.jl:418
  [4] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:1723
  [5] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:4819
  [6] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:67
  [7] compile
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:55 [inlined]
  [8] _thunk(job::GPUCompiler.CompilerJob{…}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5682
  [9] _thunk
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5680 [inlined]
 [10] cached_compilation
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5734 [inlined]
 [11] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5848
 [12] thunk_generator(world::UInt64, source::Union{…}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type, strongzero::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:6041
 [13] macro expansion
    @ ~/.julia/packages/Enzyme/9dy6G/src/rules/jitrules.jl:464 [inlined]
 [14] runtime_generic_augfwd(::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::typeof(pdf), ::Nothing, ::MultivariateTransformed{…}, ::MultivariateTransformed{…}, ::Vector{…}, ::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/9dy6G/src/rules/jitrules.jl:574
 [15] loss
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:34 [inlined]
 [16] augmented_julia_loss_22395wrap
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:0
 [17] macro expansion
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5610 [inlined]
 [18] enzyme_call
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5144 [inlined]
 [19] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/9dy6G/src/compiler.jl:5083 [inlined]
 [20] autodiff
    @ ~/.julia/packages/Enzyme/9dy6G/src/Enzyme.jl:408 [inlined]
 [21] macro expansion
    @ ~/.julia/packages/Enzyme/9dy6G/src/sugar.jl:324 [inlined]
 [22] gradient(::EnzymeCore.ReverseMode{…}, ::EnzymeCore.Const{…}, ::Vector{…})
    @ Enzyme ~/.julia/packages/Enzyme/9dy6G/src/sugar.jl:262
 [23] top-level scope
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:36
Some type information was truncated. Use `show(err)` to see complete types.

[mmikhasenko]

Thanks a lot @gdalle

Btw: I'm surprised by the struct types, Trivial{T,V<:AbstractVector{T}}, in contrary to Trivial{V<:AbstractVector}?
That T does not seem to be used elsewhere.

[gdalle]

The T can be useful if you want easy access to the element type of V, but you're right it is not essential