Hello all,
I would like to autodiff a function where some vector-valued arguments are kept constant. I thought to use Enzyme with the Const() functionality but run into issues when applying Const() to vectors.
The following is a minimal working example that reproduces the problem. I use
Enzyme v0.13.12 with julia 1.11
using Enzyme
# this will work
function f(x::Array{Float64}, c::Vector{Float64})
y = (x[1]-c[1]) * (x[1]-c[1]) + (x[2]-c[2]) * (x[2]-c[2])
return y
end;
# this won't work
function h(x::Array{Float64}, c::Vector{Float64})
y = sum( (x-c).^2 )
return y
end;
# this will work
function h2(x::Array{Float64}, c1::Float64, c2::Float64)
c = [c1, c2]
y = sum( (x-c).^2 )
return y
end;
The three functions compute the squared norm between the two vectors x and c. For example
x = [4.0, 3.0];
c = [2.0, 1.0];
f(x, c)
f(x, c) == h(x, c) == h2(x, c[1], c[2]) # returns true
Now, autodiff on f and h2 works
dx = [0.0, 0.0]
autodiff(Reverse, f, Active, Duplicated(x, dx), Const(c));
dx
2*(x-c) == dx # true
dx = [0.0, 0.0]
autodiff(Reverse, h2, Active, Duplicated(x, dx), Const(c[1]), Const(c[2]));
dx
2*(x-c) == dx # true
However, for h I get a Constant memory is stored (or returned) to a differentiable variable. error
dx = [0.0, 0.0]
autodiff(Reverse, h, Active, Duplicated(x, dx), Const(c));
dx
I am new to Julia and its AD system and puzzled by this error. It seems to me that Const(c) did not work when c is a vector? What would I need to change to make it work? Manually expanding the vector c into scalars won’t be an option for me.
Many thanks for your help.