Reverse-mode autodiff with Enzyme involving WignerSymbols

General Usage
1

One part of a function chain I’m differentiating with Enzyme involves Clebsch-Gordan coefficients generated by the WignerSymbols package. I’m able to differentiate a MWE function with Forward mode, but Reverse mode gives this error:

ERROR: LoadError: Enzyme compilation failed due to illegal type analysis.
 This usually indicates the use of a Union type, which is not fully supported with Enzyme.API.strictAliasing set to true [the default].
 Ideally, remove the union (which will also make your code faster), or try setting Enzyme.API.strictAliasing!(false) before any autodiff call.
 To toggle more information for debugging (needed for bug reports), set Enzyme.Compiler.VERBOSE_ERRORS[] = true (default false)

Caused by:
Stacktrace:
 [1] getproperty
   @ ./Base.jl:37
 [2] splitsquare
   @ ~/.julia/packages/WignerSymbols/rppM5/src/primefactorization.jl:342
 [3] _wigner3j
   @ ~/.julia/packages/WignerSymbols/rppM5/src/WignerSymbols.jl:109

Stacktrace:
  [1] julia_error(msg::String, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/errors.jl:319
  [2] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/errors.jl:210
  [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, width::Int64, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/ydGh2/src/api.jl:404
  [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::Tuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:1501
  [5] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:4436
  [6] codegen
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:3239 [inlined]
  [7] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5288
  [8] _thunk
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5288 [inlined]
  [9] cached_compilation
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5340 [inlined]
 [10] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5451
 [11] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::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)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5636
 [12] runtime_generic_augfwd(activity::Type{…}, runtimeActivity::Val{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(clebschgordan), df::Nothing, primal_1::Float64, shadow_1_1::Nothing, primal_2::Float64, shadow_2_1::Nothing, primal_3::Float64, shadow_3_1::Nothing, primal_4::Float64, shadow_4_1::Nothing, primal_5::Float64, shadow_5_1::Nothing, primal_6::Float64, shadow_6_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/rules/jitrules.jl:465
 [13] fun
    @ ~/nuclear-diffprog/MWEs/wigner.jl:11
 [14] #25
    @ ~/nuclear-diffprog/MWEs/wigner.jl:17 [inlined]
 [15] augmented_julia__25_10950wrap
    @ ~/nuclear-diffprog/MWEs/wigner.jl:0
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5218 [inlined]
 [17] enzyme_call
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:4764 [inlined]
 [18] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:4700 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:396 [inlined]
 [20] autodiff
    @ ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:524 [inlined]
 [21] macro expansion
    @ ~/.julia/packages/Enzyme/ydGh2/src/sugar.jl:326 [inlined]
 [22] gradient(::ReverseMode{false, false, FFIABI, false, false}, ::var"#25#26", ::Float64)
    @ Enzyme ~/.julia/packages/Enzyme/ydGh2/src/sugar.jl:258
 [23] top-level scope
    @ ~/nuclear-diffprog/MWEs/wigner.jl:17
 [24] include(fname::String)
    @ Base.MainInclude ./client.jl:494
 [25] top-level scope
    @ REPL[6]:1
in expression starting at /vast/home/daningburg/nuclear-diffprog/MWEs/wigner.jl:17
Some type information was truncated. Use `show(err)` to see complete types.

Strangely it does work for Forward mode, but I need Reverse for my main code. Here’s the example:

# Try to differentiate a function that includes a Clebsch-Gordan Coefficient using Enzyme
begin # Packages
    using Enzyme
    using WignerSymbols
end

j1, m1, j2, m2, j3, m3 = 3.5, -0.5, 0.0, 0.0, 3.5, -0.5
cg = clebschgordan(j1, m1, j2, m2, j3, m3)

function fun(x)
    y = clebschgordan(j1, m1, j2, m2, j3, m3) * x.^2
    return y 
end

x = 3.0
# Enzyme.gradient(Forward, x -> fun(x), x)
Enzyme.gradient(Reverse, x -> fun(x), x)
2 
>>> typeof(cg)

RationalRoots.RationalRoot{BigInt}

So this is a pretty memory flexiblly sized type which helps say meson calculations perfectly cancel where theres a symmetry. If you just want a float64 maybe look at the options here Timing Clebsch-Gordan-coefficients calls in Julia (beating c++ again) - #3 by misha_mikhasenko.

edit: or this one looks good and has specific float operations GitHub - 0382/CGcoefficient.jl: Compute CG coefficient, Racah coefficient, and Wigner 3j, 6j, 9j Symbols, and give the exact results.

1 Like
3

I was able to get CGcoefficients working with my MWE (posted below for completeness). Sadly I got a new kind of error with “core dumped” when integrating it into my main code. I’ll try some of the other options, thanks!

# Try to differentiate a function that includes a Clebsch-Gordan Coefficient using Enzyme
begin # Packages
    using Enzyme
    using WignerSymbols
    using CGcoefficient
end

function fun(x)
    # y = clebschgordan(j1, m1, j2, m2, j3, m3) * x.^2
    y = CGcoefficient.fCG(dj1, dj2, dj3, dm1, dm2, dm3) * x^2
    return y 
end

function double_hint(j)
    dj = Int(2*j)
end

function all_doubles(j1, j2, j3, m1, m2, m3)
    dj1, dj2, dj3, dm1, dm2, dm3 = double_hint(j1), double_hint(j2), double_hint(j3), double_hint(m1), double_hint(m2), double_hint(m3)
    return dj1, dj2, dj3, dm1, dm2, dm3
end

L = 14
CGcoefficient.wigner_init_float(L, "Jmax", 6)
j1, m1, j2, m2, j3, m3 = 3.5, -0.5, 1.0, 0.0, 3.5, -0.5
dj1, dj2, dj3, dm1, dm2, dm3 = all_doubles(j1, j2, j3, m1, m2, m3)
cg = clebschgordan(j1, m1, j2, m2, j3, m3)
cg2 = CGcoefficient.fCG(dj1, dj2, dj3, dm1, dm2, dm3)

x = 3.0
# Enzyme.gradient(Forward, x -> fun(x), x)
Enzyme.gradient(Reverse, x -> fun(x), x)
4

Main author of WignerSymbols.jl here. I assume that in principle you don’t want AD to go through the wigner symbol calculation, since all of its inputs are discrete numbers (half integers), there is nothing with respect to which you can differentiate. Unless I am overlooking something?

I am happy to add a package extension to tell Enzyme not to try to differentiate through these functions, but I would need some help from someone with Enzyme expertise. I guess some analogue to the ChainRulesCore.@non_differentiable macro.

2 Likes
5

I confirmed the "core dumped’ is unrelated to the use of CGcoefficients so I’m marking solved.

6

You are correct that I don’t want differentiation of the CG coefficients. I think Enzyme is smart enough to know that (whereas in my experience Zygote was not, and I used ChainRulesCore.ignore_derivative() to let it know). I think the issue is related to the rational root BigInt representation, which as I see from @cdawg is needed for certain situations but not for my own.

Perhaps something like CGcoefficients.jl’s float version of the calculations would be useful for others like me, who are doing numerical work and differentiation of functions which involve Wigner symbols.