# How can we create more basic Enzyme custom rules?

**URL:** <https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174>\
**Category:** Numerics\
**Tags:** question, autodiff, enzyme, advice\
**Created:** [March 27, 2024, 9:45am UTC](https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174 "2024-03-27T09:45:46Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![ToPo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/topo/32/205868_2.png) [@ToPo](https://discourse.julialang.org/u/ToPo)\
**Post date:** [March 27, 2024, 9:45am UTC](https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174/1 "2024-03-27T09:45:46Z")

</div>

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.

```julia
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 

```

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [March 27, 2024, 3:07pm UTC](https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174/2 "2024-03-27T15:07:53Z")

</div>

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](https://enzyme.mit.edu/index.fcgi/julia/dev/faq/#Mixed-activity:)

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.

---

<div class="post-metadata">

**Author:** ![ToPo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/topo/32/205868_2.png) [@ToPo](https://discourse.julialang.org/u/ToPo)\
**Post date:** [March 27, 2024, 3:31pm UTC](https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174/3 "2024-03-27T15:31:25Z")

</div>

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

```julia
#--------------------------------------------------------------------------------
# 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.

```julia
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}}}

```

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [March 27, 2024, 3:37pm UTC](https://discourse.julialang.org/t/how-can-we-create-more-basic-enzyme-custom-rules/112174/4 "2024-03-27T15:37:37Z")

</div>

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