Mutating array workaround for reverse AD?

beryl
Fri, 13 May 2022 11:45:45 +0000
General Usage

Hi,

I have yet another unsolved problem for reverse mode AD (or I think, this is the same problem as the one I posted previously: LoadError: ArgumentError: Converting an instance of ReverseDiff.TrackedReal{Float64, Float64, Nothing} to Float64 is not defined. Please use `ReverseDiff.value` instead). Here is the (hopefully) MWE:

using LinearAlgebra
using Zygote, ReverseDiff

function f_test_1(A, x)
    """
    u = Ax + some constant row vector
    """
    u = A*x[2:end] .+ x[1]
    return u
end

function f_test_2(A, x)
    """
    alloc inside
    """
    u = Vector{Float64}(undef, length(x)-1)
    u .= A*x[2:end] .+ x[1]
    return u
end

function f_test_3!(u, A, x)
    """
    non-alloc ver
    """
    u .= A*x[2:end] .+ x[1]
end

# derivative:
J_f_1(A, x) = Zygote.jacobian(θ -> f_test_1(A, θ), x)
J_f_2(A, x) = Zygote.jacobian(θ -> f_test_2(A, θ), x)
J_f_3(u, A, x) = Zygote.jacobian(θ -> f_test_3!(u, A, θ), x)

J_f_1 works, however _2 and _3 don’t, for example:

A = Matrix{Float64}(LinearAlgebra.I, 5, 5)
x = ones(6)
u = Vector{Float64}(undef, 5)

J_f_1(A, x) # this works
J_f_2(A, x) # throws "Mutating arrays is not supported -- called copyto!(::Vector{Float64}, _...)"
J_f_3(u,A,x) # same as _2

Also, replacing Zygote with ReverseDiff produces similar error it seems, although the error trace shows different thing (ReverseDiff gives “ArgumentError: Converting an instance of ReverseDiff.TrackedReal{Float64, Float64, Nothing} to Float64 is not defined. Please use ReverseDiff.value instead.”, I tried using ReverseDiff.value, however it results in incorrect Jacobian (all 0.)).
So, are there any workarounds for Zygote or ReverseDiff to work for mutating arrays (other than using f_test_1)? For example if I have more complicated functions which require assignments of values to the output array(s), like B[i,:,:] = … .
Thanks.

6 posts - 4 participants