LoopVectorization.jl icon indicating copy to clipboard operation
LoopVectorization.jl copied to clipboard

Hook into Zygote.jl?

Open AStupidBear opened this issue 5 years ago • 2 comments

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

AStupidBear avatar May 08 '20 15:05 AStupidBear

I would like to add support for this as well by defining a pullback.

chriselrod avatar May 08 '20 16:05 chriselrod

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)

mcabbott avatar May 09 '20 12:05 mcabbott