Fix TuringMvNormal's rand on GPU

[nmheim]

My attempt to fix #98 for TuringMvNormal by @requireing CUDA. What do you think?

[devmotion]

It seems the name randnsimilar is slightly misleading? It doesn't ensure that the element types are preserved but only if an array is created on the GPU or not.

My suggestion would be to

• make CUDA a proper dependency as recommended in its documentation:
  "Because CUDA.jl always loads, even if the user doesn't have a GPU or CUDA, you should just depend on it like any other package (and not use, e.g., Requires.jl). This ensures that breaking changes to the GPU stack will be taken into account by the package resolver when installing your package."
• do not define randnsimilar but just add dispatches such as

  function Base.rand(
      rng::Random.AbstractRNG,
      d::TuringDenseMvNormal{<:CuArray},
      n::Int...,
  )

[nmheim]

It seems the name randnsimilar is slightly misleading? It doesn't ensure that the element types are preserved but only if an array is created on the GPU or not.

good point, I'll fix the the element types

* make CUDA a proper dependency as recommended in [its documentation](https://juliagpu.gitlab.io/CUDA.jl/installation/conditional/):

As this was not an issue here before I assumed that not many people need GPU support? Adding CUDA increases loading times from about 5s to 8s on my machine. It would ofc be nicer to have it as a proper dependency, but I was not sure if that was desired.

[devmotion]

As this was not an issue here before I assumed that not many people need GPU support? Adding CUDA increases loading times from about 5s to 8s on my machine. It would ofc be nicer to have it as a proper dependency, but I was not sure if that was desired.

The increased loading times are unfortunate. In my experience using Requires tends to increase loading times as well though and, probably even worse, doesn't allow us to specify any compatibility bounds (which is annoying for the AD backends as well BTW). An alternative would be a separate package, maybe named DistributionsADCUDA or DistributionsADGPU (or the other way around), that just adds GPU support to DistributionsAD and depends on both DistributionsAD and CUDA.

[nmheim]

Ok, I think I have not understood entirely how the CPU/GPU rng is used (which was the reason for not using randn! and introducing CUDA as a dependency). If I run:

using DistributionsAD
using CUDA
using Random
using Zygote

Random.seed!(0)

μ = rand(Float32,2) |> cu
σ = rand(Float32,2) |> cu
d = DistributionsAD.TuringMvNormal(μ, σ)
display(rand(d))

f(μ,σ) = sum(rand(DistributionsAD.TuringMvNormal(μ,σ)))

for a in Zygote.gradient(f, μ, σ) display(a) end

with the code from the commit above, fixing only the cpu seed fixes the output of rand(d). what am I missing here? the way it is written now would get by without the CUDA dependency.

fixing only the cpu seed and using CUDA.randn would ofc not fix the output of rand(d), but is that a problem?

DistributionsAD.jl/src/multivariate.jl

Line 101 in 65c362d

Distributions.rand(d::TuringDenseMvNormal, n::Int...) = rand(Random.GLOBAL_RNG, d, n...)

seems to use cpu rng...?

[devmotion]

I don't remember exactly the original issue - was it caused mainly by rand not preserving the element types (Float32 vs Float64)?

[nmheim]

the original issue (#98) was that randn inside rand(::MvNormal) returned Arrays instead of CuArrays when the distributions where on the GPU. Ah, and I think I see the problem. like this the randn! calls are forwarded to scalar operations. sorry, I forgot about that

[nmheim]

Ok, so like this randnsimilar should create arrays with preserved eltype and array type.
Introduces dependency on CUDA (and Adapt, which is already part of CUDA), which needs julia 1.4 though...

[devmotion]

Introduces dependency on CUDA (and Adapt, which is already part of CUDA), which needs julia 1.4 though...

The latest release of CUDA even requires Julia 1.5. However, IMO we shouldn't drop support for Julia < 1.5 currently.

[nmheim]

The latest release of CUDA even requires Julia 1.5. However, IMO we shouldn't drop support for Julia < 1.5 currently.

Yes, agreed, the commit above restricts CUDA to <1.3.3 which is still compatible with julia 1.3

[nmheim]

I think the tests are currently failing because of JuliaDiff/ChainRules.jl#262. So I think once JuliaDiff/ChainRules.jl#263 is merged we should be fine.

Edit: acutally, it doesn't seem like it. the gradient test tracebacks are a bit tricky to read ^^
The gradients through rand are not tested currently right?

[devmotion]

The gradients through rand are not tested currently right?

No, we only test gradients of logpdf and loglikelihood. rand calls are not differentiable in ChainRules anyways, see e.g. https://github.com/JuliaDiff/ChainRules.jl/blob/master/src/rulesets/Random/random.jl.

Edit: Just saw your PR, yes, I guess it would make sense to test gradients there as well if it is supported by some AD backend.

[nmheim]

I guess it would make sense to test gradients there as well if it is supported by some AD backend.

I'll open an issue for this

[nmheim]

I am still a bit puzzled by the failing TuringPoissonBinomial though. could that also be related to some new @non_differentiable rules?

[devmotion]

Possibly. Tests on master passed with ChainRules v0.7.14 and ChainRulesCore v0.9.6, this PR fails with ChainRules v0.7.17 and ChainRulesCore v0.9.7. The differences are JuliaDiff/ChainRulesCore.jl@v0.9.6...v0.9.7 and JuliaDiff/ChainRules.jl@v0.7.14...v0.7.17.

[devmotion]

I guess the problem might be the call of isapprox in

DistributionsAD.jl/src/univariate.jl

Line 43 in f754a19

check_args && @assert all(x -> x >= -ϵ, pb) && isapprox(sum(pb), 1, atol = ϵ)

and the definition in
https://github.com/JuliaDiff/ChainRules.jl/blob/0baf7bba2a9a2f235bf2a85edfdc6209f8cf137d/src/rulesets/Base/nondiff.jl#L123.

[nmheim]

I guess the problem might be the call of isapprox in

DistributionsAD.jl/src/univariate.jl

Line 43 in f754a19

check_args && @assert all(x -> x >= -ϵ, pb) && isapprox(sum(pb), 1, atol = ϵ)

and the definition in JuliaDiff/ChainRules.jl@v0.7.14...v0.7.17diff-608aeb35114f6c76d5282c97cba14d6aR123.

Yes, you are right, thats it. The latest commit contains a hacky fix, because @non_differentiable cannot deal with keyword args. What do you think?

[devmotion]

LGTM, I just have the following comments left:

• Can you please bump the version number to 0.6.8 so we can make a new release with the bugfix and new feature? (Ideally, they should be separated but IMO it's OK to keep it in one PR now)
• Can you open an issue over at ChainRulesCore to inform them about the breakage of isapprox and raise awareness for keyword argument support in @non_differentiable?
• Can you open an issue in DistributionsAD that we should remove _isapprox as soon as the underlying problem is fixed upstream? Otherwise we might just not remember it.

[devmotion]

Great, LGTM! Thanks for the PR!

It seems to work but we should make sure to actually run some GPU tests in the future as well: https://github.com/JuliaGPU/gitlab-ci