Julia :Zygote.@adjoint 来自 Enzyme.autodiff

Question

给定下面的函数 f!:

function f!(s::Vector, a::Vector, b::Vector)
  
  s .= a .+ b
  return nothing

end # f!

如何根据

为 Zygote定义伴随物

Enzyme.autodiff(f!, Const, Duplicated(s, dz_ds)。Duplicated(a, zero(a)), Duplicated(b, zero(b))) ?

Zygote.@adjoint f!(s, a, b) = f!(s, a, b), # What would come here ?

Answer 1

想出了一个办法，在这里分享。

对于给定函数 foo，Zygote.pullback(foo, args...) 返回 foo(args...) 并且向后传递(允许梯度计算)。

我的目标是告诉 Zygote 使用 Enzyme 进行反向传递。

这可以通过 Zygote.@adjoint 完成(参见更多 here )。

在数组值函数的情况下，Enzyme 需要一个返回nothing 并且其结果在args 中的变异版本(查看更多here ).

问题帖中的函数 f! 是两个数组之和的 Enzyme 兼容版本。

因为 f! 返回 nothing，Zygote 只会返回 nothing当对传递给我们的某些梯度调用反向传递时。

一个解决方案是将 f! 放在返回数组 s 的包装器(比如 f)中

并为 f 定义 Zygote.@adjoint，而不是 f!。

因此，

function f(a::Vector, b::Vector)

  s = zero(a)
  f!(s, a, b)
  return s

end

function enzyme_back(dzds, a, b)

  s    = zero(a)
  dzda = zero(dzds)
  dzdb = zero(dzds)
  Enzyme.autodiff(
    f!,
    Const,
    Duplicated(s, dzds),
    Duplicated(a, dzda),
    Duplicated(b, dzdb)
  )
  return (dzda, dzdb)

end

和

Zygote.@adjoint f(a, b) = f(a, b), dzds -> enzyme_back(dzds, a, b)

通知 Zygote 在反向传递中使用 Enzyme。

---

最后，您可以检查调用 Zygote.gradient 或者在

g1(a::Vector, b::Vector) = sum(abs2, a + b)

或

g2(a::Vector, b::Vector) = sum(abs2, f(a, b))

产生相同的结果。