refactoring

Author: sunxd3

ext/NormalizingFlowsCUDAExt.jl

@@ -4,22 +4,22 @@ using CUDA
 using NormalizingFlows
 using NormalizingFlows: Bijectors, Distributions, Random
 
-function NormalizingFlows.rand_device(
+function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
-    return rand_cuda(rng, s)
+    return _cuda_rand(rng, s)
 end
 
-function NormalizingFlows.rand_device(
+function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
 )
-    return rand_cuda(rng, s, n)
+    return _cuda_rand(rng, s, n)
 end
 
-function rand_cuda(
+function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
@@ -28,7 +28,7 @@ function rand_cuda(
     )
 end
 
-function rand_cuda(
+function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
@@ -39,31 +39,31 @@ function rand_cuda(
 end
 
 # ! this is type piracy
-# replace scalar indexing
+# replacing original function with scalar indexing
 function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVecOrMat)
     Random.randn!(rng, x)
     Distributions.unwhiten!(d.Σ, x)
     x .+= d.μ
     return x
 end
 
-# to enable `rand_device(rng:CUDA.RNG, flow[, num_samples])`
-function NormalizingFlows.rand_device(rng::CUDA.RNG, td::Bijectors.TransformedDistribution)
-    return rand_cuda(rng, td)
+# to enable `_device_specific_rand(rng:CUDA.RNG, flow[, num_samples])`
+function NormalizingFlows._device_specific_rand(rng::CUDA.RNG, td::Bijectors.TransformedDistribution)
+    return _cuda_rand(rng, td)
 end
 
-function NormalizingFlows.rand_device(
+function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG, td::Bijectors.TransformedDistribution, num_samples::Int
 )
-    return rand_cuda(rng, td, num_samples)
+    return _cuda_rand(rng, td, num_samples)
 end
 
-function rand_cuda(rng::CUDA.RNG, td::Bijectors.TransformedDistribution)
-    return td.transform(rand_cuda(rng, td.dist))
+function _cuda_rand(rng::CUDA.RNG, td::Bijectors.TransformedDistribution)
+    return td.transform(_cuda_rand(rng, td.dist))
 end
 
-function rand_cuda(rng::CUDA.RNG, td::Bijectors.TransformedDistribution, num_samples::Int)
-    samples = rand_cuda(rng, td.dist, num_samples)
+function _cuda_rand(rng::CUDA.RNG, td::Bijectors.TransformedDistribution, num_samples::Int)
+    samples = _cuda_rand(rng, td.dist, num_samples)
     res = reduce(
         hcat,
         map(axes(samples, 2)) do i

src/NormalizingFlows.jl

@@ -83,10 +83,44 @@ function train_flow(
 end
 
 include("optimize.jl")
-include("objectives.jl")
 
-function rand_device end
+# objectives
+include("objectives/elbo.jl")
+include("objectives/loglikelihood.jl") # not fully tested
 
-include("rand_device.jl")
+"""
+    _device_specific_rand
+
+By default dispatch to `Random.rand`, but maybe overload when the random number 
+generator is device specific (e.g. `CUDA.RNG`).
+"""
+function _device_specific_rand end
+
+function _device_specific_rand(
+    rng::Random.AbstractRNG,
+    s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
+)
+    return Random.rand(rng, s)
+end
+
+function _device_specific_rand(
+    rng::Random.AbstractRNG,
+    s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
+    n::Int,
+)
+    return Random.rand(rng, s, n)
+end
+
+function _device_specific_rand(
+    rng::Random.AbstractRNG, td::Bijectors.TransformedDistribution
+)
+    return Random.rand(rng, td)
+end
+
+function _device_specific_rand(
+    rng::Random.AbstractRNG, td::Bijectors.TransformedDistribution, n::Int
+)
+    return Random.rand(rng, td, n)
+end
 
 end

src/objectives.jl

@@ -33,11 +33,11 @@ function elbo(flow::Bijectors.MultivariateTransformed, logp, xs::AbstractMatrix)
 end
 
 function elbo(rng::AbstractRNG, flow::Bijectors.MultivariateTransformed, logp, n_samples)
-    return elbo(flow, logp, rand_device(rng, flow.dist, n_samples))
+    return elbo(flow, logp, _device_specific_rand(rng, flow.dist, n_samples))
 end
 
 function elbo(rng::AbstractRNG, flow::Bijectors.UnivariateTransformed, logp, n_samples)
-    return elbo(flow, logp, rand_device(rng, flow.dist, n_samples))
+    return elbo(flow, logp, _device_specific_rand(rng, flow.dist, n_samples))
 end
 
 function elbo(flow::Bijectors.TransformedDistribution, logp, n_samples)

src/objectives/elbo.jl

@@ -35,8 +35,8 @@ using Bijectors, CUDA, Distributions, Flux, LinearAlgebra, Test
 
     @testset "$dist" for dist in dists
         CUDA.allowscalar(true)
-        x = NormalizingFlows.rand_device(CUDA.default_rng(), dist)
-        xs = NormalizingFlows.rand_device(CUDA.default_rng(), dist, 100)
+        x = NormalizingFlows._device_specific_rand(CUDA.default_rng(), dist)
+        xs = NormalizingFlows._device_specific_rand(CUDA.default_rng(), dist, 100)
         @test_nowarn logpdf(dist, x)
         @test x isa CuArray
         @test xs isa CuArray
@@ -52,8 +52,8 @@ using Bijectors, CUDA, Distributions, Flux, LinearAlgebra, Test
         )
         flow = Bijectors.transformed(dist, ComposedFunction(pl1, pl2))
 
-        y = NormalizingFlows.rand_device(CUDA.default_rng(), flow)
-        ys = NormalizingFlows.rand_device(CUDA.default_rng(), flow, 100)
+        y = NormalizingFlows._device_specific_rand(CUDA.default_rng(), flow)
+        ys = NormalizingFlows._device_specific_rand(CUDA.default_rng(), flow, 100)
         @test_nowarn logpdf(flow, y)
         @test y isa CuArray
         @test ys isa CuArray