Vectorization of assume

[xukai92]

The most easy one is to vectorize observe where the distributions are all the same. observe for distributions which are different and assume are challenging. I will do the easiest one first.

-- update on 30 May --

TODOs

• support Vector-variate (interface)
• support Matrix-variate (interface)
• always change things in-place for vectorised assume()

Left todos are moved to PR #281

[xukai92]

I did a simple check:

# Test for vectorize UnivariateDistribution
@model vdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  x ~ Normal(m, sqrt(s))
  # for i = 1:length(x)
  #   x[i] ~ Normal(m, sqrt(s))
  # end
  return s, m
end

with x=randn(1000).

With the loop way, it will take 1m26s but with vectorize version (bb25da0) only 2.25 seconds.

If we can also do that for at lease iid distributions priors, we can certainly speed up the model() quite a lot.

[xukai92]

Cont'd run by NUTS(2500, 500, 0.65)

[yebai]

With the loop way, it will take 1m26s but with vectorize version (bb25da0) only 2.25 seconds.
If we can also do that for at lease iid distributions priors, we can certainly speed up the model() quite a lot.

That's exciting! We can slowly add more vectorisation support.

[xukai92]

One possible and simple way to vectorize i.i.d priors is to create some customized distributions like MvDirichlet and MvCategorical, with original parameter and number of variables provided, then we can then call the original with rand(d, n) to assign them to the left. I guess VarInfo can just treat them as normal variables?

[xukai92]

Vec*** would be a better name without overriding any thing in the distribution package. We can simply generate Vec*** for all distributions and have rand and logpdf for them by calling the corresponding ones by rand(d, n) and logpdf(d, n).

[xukai92]

I've done the vectorization for assume with univarite distributions. Here are some experiment results

N = 100
setchunksize(N)
alg = HMC(1000, 0.2, 4)

@model vdemo() = begin
  x = Vector{Real}(N)
  for i = 1:N
    x[i] ~ Normal(0, sqrt(4))
  end
end

res = sample(vdemo(), alg)
println(mean(mean(res[:x])))



# Test for vectorize UnivariateDistribution
@model vdemo() = begin
  x = Vector{Real}(N)
  x ~ [Normal(0, 2)]
end


res = sample(vdemo(), alg)
println(mean(mean(res[:x])))


@model vdemo() = begin
  x ~ MvNormal(zeros(N), 2 * ones(N))
end

res = sample(vdemo(), alg)
println(mean(mean(res[:x])))

Time for each

• loop: 18.15s
• vec: 14.32s (current implementation of af2b6c6)
• MvNormal: 10.12s

I believe the gap between vec and Mv means there is still space to improve Turing in terms of vectorization.

[yebai]

• Univariate
• Vector-variate
• Matrix-variate

[xukai92]

NOTE

• We can recover dimension by length_of_vec_in_vi / length_of_dist
• Always change things in-place for vectorised assume()
• TODOs are added to the top so that GitHub can track it

[xukai92]

May worth have a look at if @simd can speed the fake vec of matrix variate implemeatnion in Distributions.jl
Ref:

[xukai92]

Closed by #281