Different api for scalar input in Enzyme.jl

I’m puzzled by a difference in the API of Reverse mode with functions of scalar inputs versus those of other (e.g. vector inputs). It seems to be the case that for a scalar function of a scalar input, you must use Reverse mode with Active arguments, or wrap them as an array.

Take this simple example. It seems like the correct way to use Duplicated arguments for a scalar inputs is to wrap them in an array structure as I’ve done in the fourth call below.

using Enzyme

g(x) = sum(x.^2)

y = 1.;
dy = 0.;
@show autodiff(ReverseWithPrimal, g, Active, Duplicated(y,dy))
@show y,dy

y = [1., 2.];
dy = [0.,0.];
@show autodiff(ReverseWithPrimal, g, Active, Duplicated(y,dy))
@show y,dy

y = 1.;
@show dy=autodiff(Reverse, g, Active, Active(y))[1][1]
@show y,dy

y = [1.];
dy = [0.];
@show autodiff(ReverseWithPrimal, g, Active, Duplicated(y,dy))
@show y,dy

Output:

autodiff(ReverseWithPrimal, g, Active, Duplicated(y, dy)) = ((nothing,), 1.0)
(y, dy) = (1.0, 0.0)
autodiff(ReverseWithPrimal, g, Active, Duplicated(y, dy)) = ((nothing,), 5.0)
(y, dy) = ([1.0, 2.0], [2.0, 4.0])
dy = ((autodiff(Reverse, g, Active, Active(y)))[1])[1] = 2.0
(y, dy) = (1.0, 2.0)
autodiff(ReverseWithPrimal, g, Active, Duplicated(y, dy)) = ((nothing,), 1.0)
(y, dy) = ([1.0], [2.0])

Is there a design reason why Duplicated won’t work with scalars?

1 Like

Scalars are passed by value to Duplicated whereas arrays are passed by reference. In order to modify the underlying data during autodiff, it needs to be passed by reference.

Got it. Thanks!