How to port NiLang to ChainRules
In How to port NiLang to Zygote we showed the way to insert Nilang-based gradient as Zygote's pullback/adjoint. Given that ChainRules is now the core of many AD packages including Zygote, extending ChainRules.rrule
with Nilang does the same job, except that it affects all ChainRules-based AD packages and not just Zygote.
We'll use the same example as How to port NiLang to Zygote, so you might need to restart your Julia to get a fresh environment.
using NiLang, NiLang.AD, Zygote, ChainRules
Let's start from the Julia native implementation of norm2
function.
function norm2(x::AbstractArray{T}) where T
out = zero(T)
for i=1:length(x)
@inbounds out += x[i]^2
end
return out
end
norm2 (generic function with 1 method)
Zygote is able to generate correct dual function, i.e., gradients, but much slower than the primal function norm2
using BenchmarkTools
x = randn(1000);
original_grad = norm2'(x)
@benchmark norm2'($x) seconds=1
BenchmarkTools.Trial: 296 samples with 1 evaluation.
Range (min … max): 2.359 ms … 11.974 ms ┊ GC (min … max): 0.00% … 51.01%
Time (median): 2.488 ms ┊ GC (median): 0.00%
Time (mean ± σ): 3.401 ms ± 2.406 ms ┊ GC (mean ± σ): 22.27% ± 21.43%
█▅
██▄▁▁▄▁▄▆▁▄▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄█▅▆▅▆▅▁▁▁▁▁▁▁▁▁▁▁▄▅▇ ▅
2.36 ms Histogram: log(frequency) by time 11.5 ms <
Memory estimate: 8.36 MiB, allocs estimate: 19059.
The primal function is
@benchmark norm2($x) seconds=1
BenchmarkTools.Trial: 10000 samples with 10 evaluations.
Range (min … max): 1.300 μs … 2.610 μs ┊ GC (min … max): 0.00% … 0.00%
Time (median): 1.310 μs ┊ GC (median): 0.00%
Time (mean ± σ): 1.316 μs ± 43.614 ns ┊ GC (mean ± σ): 0.00% ± 0.00%
█
▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃ ▂
1.3 μs Histogram: frequency by time 1.33 μs <
Memory estimate: 0 bytes, allocs estimate: 0.
Then we have the reversible implementation
@i function r_norm2(out::T, x::AbstractArray{T}) where T
for i=1:length(x)
@inbounds out += x[i]^2
end
end
The gradient generated by NiLang is much faster, which is comparable to the forward program
@benchmark (~r_norm2)(GVar($(norm2(x)), 1.0), $(GVar(x))) seconds=1
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Range (min … max): 40.101 μs … 74.301 μs ┊ GC (min … max): 0.00% … 0.00%
Time (median): 41.300 μs ┊ GC (median): 0.00%
Time (mean ± σ): 41.376 μs ± 1.019 μs ┊ GC (mean ± σ): 0.00% ± 0.00%
▂▅▆▇█▇▇▆▄▁
▂▂▃▄▅▇██████████▇▆▅▅▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▁▂▁▂▁▁▁▂▁▁▁▁▁▁▁▁▁▁▁▁▂ ▃
40.1 μs Histogram: frequency by time 46.3 μs <
Memory estimate: 0 bytes, allocs estimate: 0.
By defining our custom rrule
using Nilang's gradient implementation, Zygote
automaticallly gets boosted because it internally uses the available ChainRules ruleset. Here we need to create a new symbol here because otherwise Zygote will still use the previously generated slow implementation.
norm2_faster(x) = norm2(x)
function ChainRules.rrule(::typeof(norm2_faster), x::AbstractArray{T}) where T
out = norm2_faster(x)
function pullback(ȳ)
ChainRules.NoTangent(), grad((~r_norm2)(GVar(out, ȳ), GVar(x))[2])
end
out, pullback
end
@assert norm2_faster'(x) ≈ original_grad
See, much faster
@benchmark norm2_faster'(x) seconds=1
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Range (min … max): 43.800 μs … 5.332 ms ┊ GC (min … max): 0.00% … 98.24%
Time (median): 47.101 μs ┊ GC (median): 0.00%
Time (mean ± σ): 49.094 μs ± 74.491 μs ┊ GC (mean ± σ): 2.13% ± 1.39%
▅▂▇█▃▅ ▂
▂▂▂▃▃▅▅██████████▆▆▄▄▃▃▃▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▂▂▂▃▃▃▃▃▃▃▃▃▃▃▂▂▂ ▃
43.8 μs Histogram: frequency by time 58.8 μs <
Memory estimate: 23.69 KiB, allocs estimate: 2.
This page was generated using Literate.jl.