How can we create more basic Enzyme custom rules?

Specific Domains Numerics
, , ,
1

I am trying to create custom rules for Enzyme.
I looked at the examples on Enzyme’s official website, but unfortunately could not fully understand them due to difficulty.

As a very simple example to try, I created a program with a custom rule for matrix products, but it also fails.
I need advice on how to solve the following program’s bad points.

I think a few simple examples would help everyone understand.

using Enzyme
import .EnzymeRules: augmented_primal, reverse
using .EnzymeRules
#--------------------------------------------------------------------------------
# Function to which we want to apply custom rules
# matrix-matrix product
#--------------------------------------------------------------------------------
function g(A::Matrix{Float64}, B::Matrix{Float64})
    return A * B
end
#--------------------------------------------------------------------------------
# Enzyme custom rules
#--------------------------------------------------------------------------------
function EnzymeRules.augmented_primal(config::ConfigWidth{1}, func::Const{typeof(g)}, ::Type{<:Active}, A::Duplicated, B::Duplicated)
    println("In custom augmented primal rule.")
    # Compute primal
    primal = func.val(A.val, B.val)
    # Return an AugmentedReturn object
    return AugmentedReturn(primal, A.val, B.val)
end
#--------------------------------------------------------------------------------
# Enzyme custom rules
#--------------------------------------------------------------------------------
function EnzymeRules.reverse(config::ConfigWidth{1}, func::Const{typeof(g)}, dC::Active, A::Duplicated, B::Duplicated)
    println("In custom reverse rule.")
    # dA
    A.dval = dC.val * B'.val
    # dB
    B.dval = A'.val * dC.val
    #
    return (nothing, nothing)
end
#--------------------------------------------------------------------------------
# Function to be differentiated
#--------------------------------------------------------------------------------
function eval(x::Vector{Float64})
    n = 3
    A = Matrix{Float64}(undef, n, n)
    for i in 1 : n
        for j in 1 : n
            A[i, j] = 2.0 * x[i] - 3.0 * x[j]
        end
    end
    B = Matrix{Float64}(undef, n, n)
    for i in 1 : n
        for j in 1 : n
            B[i, j] = 2.0 * i - 3.0 * j
        end
    end

    # matrix-matrix product
    C = g(A, B)

    return C[1, 1] + C[2, 2]
end
#--------------------------------------------------------------------------------
# main
#--------------------------------------------------------------------------------
x  = [3.0, 1.0, 2.0]
dx = [0.0, 0.0, 0.0]

# compute gradient
autodiff(Reverse, eval, Duplicated(x, dx))
@show dx
2

Here you’ve marked the return of the function you’re writing a custom rule for as active (the Type{<:Active}). However, since it’s returning a matrix, it can never be active (only duplicated or const).

These docs may be helpful: FAQ · Enzyme.jl

Active variables are used for immutable variables (like Float64), whereas Duplicated variables are used for mutable variables (like Vector{Float64}). Speciically, since Active variables are immutable, functions with Active inputs will return the adjoint of that variable. In contrast Duplicated variables will have their derivatives +='d in place.

3

Thank you!
I rewrote the following and “augmented_primal()” is now recognized.

#--------------------------------------------------------------------------------
# Enzyme custom rules
#--------------------------------------------------------------------------------
function EnzymeRules.augmented_primal(config::ConfigWidth{1}, func::Const{typeof(g)}, ::Type{<:Duplicated}, A::Duplicated, B::Duplicated)
    println("In custom augmented primal rule.")
    # Compute primal
    primal = func.val(A.val, B.val)
    # Return an AugmentedReturn object
    return AugmentedReturn(primal, A.val, B.val)
end

Next, the following error message appears and “reverse()” is not recognized.
Is there a mistake in the type of the dC argument?
Is this also A? I tried but could not calculate it properly.

ERROR: Enzyme execution failed.
Enzyme: No custom reverse rule was applicable for Tuple{ConfigWidth{1, true, true, (false, true, true)}, Const{typeof(g)}, Type{Duplicated{Matrix{Float64}}}, Matrix{Float64}, Duplicated{Matrix{Float64}}, Duplicated{Matrix{Float64}}}
4

Similarly you need to update your reverse pass rule to specify the return type as Type{<:Duplicated} as well