LoopVectorization.jl
LoopVectorization.jl copied to clipboard
Hook into Zygote.jl?
I want to take gradient of a vectorized expression, however the adjoint of vmaterialize is not defined.
using Zygote
using LoopVectorization
foo(x) = sum(@avx tanh.(x))
x = rand(10, 10)
foo'(x)
MethodError: no method matching vmaterialize(::Array{Float64,2}, ::Val{:Main})
Closest candidates are:
vmaterialize(::Base.Broadcast.Broadcasted, ::Val{Mod}) where Mod at /BIGDATA1/highchain_ylu_1/shared/.julia/packages/LoopVectorization/MZ8Kb/src/broadcast.jl:283
I would like to add support for this as well by defining a pullback.
Perhaps of interest:
using Tullio
bar(x) = sum(@tullio y[i,j] := tanh(x[i,j]))
bar'(x)
bar2(x) = @tullio s := tanh(x[i,j]) # without allocating y
bar2'(x)