可以想出一种方法,这里分享一下。
对于给定的函数
foo
,
Zygote.pullback(foo, args...)
返回
foo(args...)
和反向传递(允许梯度计算)。
我的目标是告诉
Zygote
在反向传递中使用
Enzyme
。
可以通过
Zygote.@adjoint
来实现这一点(详情
见此处)。
对于数组值函数,
Enzyme
需要一个返回
nothing
和其结果在
args
中的可变版本(详情
见此处)。
问题帖子中的函数
f!
是两个数组和的
Enzyme
兼容版本。
由于
f!
返回
nothing
,所以当我们调用反向传递时,
Zygote
会简单地返回
nothing
。
解决方案是将
f!
放在一个包装器中(比如
f
),使其返回数组
s
,
并为
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))
产生相同的结果。