# Enzyme custom rule alloactions

**URL:** <https://discourse.julialang.org/t/enzyme-custom-rule-alloactions/128491>\
**Category:** Performance\
**Tags:** autodiff, allocations, enzyme\
**Created:** [April 28, 2025, 2:11pm UTC](https://discourse.julialang.org/t/enzyme-custom-rule-alloactions/128491 "2025-04-28T14:11:20Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![c\_sell](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/c_sell/32/38854_2.png) [@c\_sell](https://discourse.julialang.org/u/c_sell)\
**Post date:** [April 28, 2025, 2:11pm UTC](https://discourse.julialang.org/t/enzyme-custom-rule-alloactions/128491/1 "2025-04-28T14:11:20Z")

</div>

Hello,  
I got a problem with allocations when using custom rules in Enzyme.  
The example below is just a toy example but the same problem arises there.  
In total I get 4 alloactions.  
2 by the line `primal = func.val(y.val, x.val)` and  
1 by `x.dval .+= 2 .* xval .* dret.val` and 1 by `x.dval .+= 2 .* xval .* y.dval`.  
If I dont define any custom rule I measure 0 allocations.

Code:

```julia
using Enzyme
import Enzyme.EnzymeRules: forward, reverse, augmented_primal
using Enzyme.EnzymeRules

x = [3.0, 1.0]
dx = [1.0, 0.0]
y = [0.0, 0.0]
dy = [0.0, 0.0]
function f(y, x)
    y .= x.^2
    return sum(y)
end

gf(y, x) = f(y, x)^2 # function to differentiate

function EnzymeRules.augmented_primal(config::R, 
    func::Const{F}, ::Type{A},
    y::Duplicated{Vector{Float64}}, x::Duplicated{Vector{Float64}}) where {R<:RevConfigWidth{1},F <: typeof(f),A<:Active}
    primal = func.val(y.val, x.val)
    return AugmentedReturn(primal, nothing, nothing)
end 

function EnzymeRules.reverse(config::RevConfigWidth{1}, func::Const{F}, dret::Active, tape,
    y::Duplicated, x::Duplicated) where F <: typeof(f)
    # println("In custom reverse rule.")
    # retrieve x value, either from original x or from tape if x may have been overwritten.
    xval = x.val
    # accumulate dret into x's shadow. don't assign!
    x.dval .+= 2 .* xval .* dret.val
    # # also accumulate any derivative in y's shadow into x's shadow.
    x.dval .+= 2 .* xval .* y.dval
    make_zero!(y.dval)
    return (nothing, nothing)
end

Enzyme.autodiff(Enzyme.Reverse, gf, Enzyme.Duplicated(y, dy), Enzyme.Duplicated(x, dx))

using Chairmarks
@b Enzyme.autodiff(Enzyme.Reverse, $gf, Enzyme.Duplicated($y, $dy), Enzyme.Duplicated($x, $dx))

```

Thank you for you help