LoopVectorization.jl/test/forwarddiffext.jl at db7af8ca6a737f9bd3711a77b23bdbb33caf69d4 · mxpoch/LoopVectorization.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
using Base: Forward

using NNlib, LoopVectorization, VectorizationBase, ForwardDiff, Test
randnvec() = Vec(ntuple(_ -> randn(), pick_vector_width(Float64))...)

tovec(x::Vec{W,T}) where {W,T} = T[Tuple(x)...]
tovec(x::VecUnroll) = reduce(vcat, map(tovec, VectorizationBase.data(x)))
function tovec(x::ForwardDiff.Dual{T,V,N}) where {T,V,N}
  v = tovec(ForwardDiff.value(x))
  dv = map(tovec, Tuple(ForwardDiff.partials(x)))
  D = ForwardDiff.Dual{T,eltype(v),N}
  ret = Vector{D}(undef, length(v))
  for i in eachindex(v)
    ret[i] = ForwardDiff.Dual(v[i], map(Base.Fix2(Base.getindex, i), dv)...)
  end
  return ret
end

vx0 = randnvec()
vx1 = randnvec()
vx2 = randnvec()
vx3 = randnvec()
vx4 = randnvec()
vx5 = randnvec()

vd0 = ForwardDiff.Dual(vx0, vx1, vx2, vx3, vx4, vx5)

vu0 = VecUnroll((vx0, vx1))
vu1 = VecUnroll((vx2, vx3))
vu2 = VecUnroll((vx4, vx5))

vud = ForwardDiff.Dual(vu0, vu1, vu2)

@test reinterpret(Float64, tovec(NNlib.leakyrelu(vd0))) ≈
      reinterpret(Float64, NNlib.leakyrelu.(tovec(vd0)))
@test reinterpret(Float64, tovec(NNlib.leakyrelu(vud))) ≈
      reinterpret(Float64, NNlib.leakyrelu.(tovec(vud)))

@test reinterpret(Float64, tovec(NNlib.relu(vd0))) ≈
      reinterpret(Float64, NNlib.relu.(tovec(vd0)))
@test reinterpret(Float64, tovec(NNlib.relu(vud))) ≈
      reinterpret(Float64, NNlib.relu.(tovec(vud)))