julia - 如何看Zygote差异化功能实现？

Question

我在 .jl 文件中编写了一个简单的函数，我可以使用 forward 成功区分它。但是我是 Julia 的新手，我不明白如何查看为微分函数生成的源代码。我已经尝试了各种各样的东西，比如 @code_lowered Zygote.forward(maxPool, [1.0, 2.0]) 和 @code_lowered Zygote.forward(maxPool) 但它们只是显示我转发自己的电话。

如何查看 Zygote 为正向和反向传递生成的代码？

using Pkg
using Zygote, ForwardDiff

function size1d(v)
    return size(v)[1]
end

function max(a, b)
    if a > b
        a
    else
        b
    end
end

function maxPool(v)
    return [max(v[2 * i - 1], v[2 * i])
            for i in 1:div(size1d(v), 2)]
end

v = [1.0, 2.0, 3.0, 4.0]
df = [20.0, 30.0]

println("maxPool(v):")
println(maxPool(v))
println()

println("maxAdjoint:")
maxAdjoint = Zygote.forward(max, 3.0, 4.0)[2]
println(maxAdjoint(1.0))
println()

println("maxPoolAdjoint:")
maxPoolAdjoint = Zygote.forward(maxPool, v)[2]
println(maxPoolAdjoint(df))

Answer 1

Zygote 有自己的宏 Zygote.@code_adjoint 用于显示降低的 adjoint 代码，即在反向模式下生成函数梯度的代码。不过我不确定转发模式。

这是一个反向模式的简单例子:

julia> using Zygote

julia> f(x) = 2x + 1
f (generic function with 1 method)

julia> @code_lowered f(1)
CodeInfo(
1 ─ %1 = 2 * x
│   %2 = %1 + 1
└──      return %2
)

julia> Zygote.@code_adjoint f(1)
Zygote.Adjoint(1: (%3, %4 :: Zygote.Context, %1, %2)
  %5 = Zygote._forward(%4, Main.:*, 2, %2)
  %6 = Base.getindex(%5, 1)
  %7 = Base.getindex(%5, 2)
  %8 = Zygote._forward(%4, Main.:+, %6, 1)
  %9 = Base.getindex(%8, 1)
  %10 = Base.getindex(%8, 2)
  return %9
, 1: (%1)
  %2 = (@10)(%1)
  %3 = Zygote.gradindex(%2, 2)
  %4 = (@7)(%3)
  %5 = Zygote.gradindex(%4, 3)
  %6 = Zygote.tuple(nothing, %5)
  return %6
)

我们可能会担心这个降低的伴随代码的长度和明显的复杂性，梯度很慢，但我们可以检查 LLVM 代码以确保所有内容最终都被删除:

julia> @code_llvm f'(1)

;  @ /Users/mason/.julia/packages/Zygote/SAZMM/src/compiler/interface.jl:50 within `#34'
define i64 @"julia_#34_18250"(i64) {
top:
  ret i64 2
}