clean up

Author: sunxd3

ext/NormalizingFlowsCUDAExt.jl

@@ -2,14 +2,12 @@ module NormalizingFlowsCUDAExt
 
 using CUDA
 using NormalizingFlows
-using NormalizingFlows: Random, Distributions, Bijectors
+using NormalizingFlows: Bijectors, Distributions, Random
 
-# to enable `rand_device(rng:CUDA.RNG, dist[, num_samples])`
 function NormalizingFlows.rand_device(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
-    println("gpu rand")
     return rand_cuda(rng, s)
 end
 
@@ -18,7 +16,6 @@ function NormalizingFlows.rand_device(
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
 )
-    println("gpu rand")
     return rand_cuda(rng, s, n)
 end
 
@@ -41,10 +38,9 @@ function rand_cuda(
     )
 end
 
-# Question: is this type piracy okay? 
-# (it's probably not ideal but this is sensible enough for now )
+# ! this is type piracy
+# replace scalar indexing
 function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVecOrMat)
-    # Replaced usage of scalar indexing.
     Random.randn!(rng, x)
     Distributions.unwhiten!(d.Σ, x)
     x .+= d.μ

src/NormalizingFlows.jl

@@ -1,18 +1,19 @@
 module NormalizingFlows
 
+using ADTypes
 using Bijectors
+using Distributions
+using LinearAlgebra
 using Optimisers
-using LinearAlgebra, Random, Distributions, StatsBase
 using ProgressMeter
-using ADTypes
+using Random
+using StatsBase
 import DifferentiationInterface as DI
 
 using DocStringExtensions
 
 export train_flow, elbo, loglikelihood
 
-export rand_device
-
 """
     train_flow([rng::AbstractRNG, ]vo, flow, args...; kwargs...)
 
@@ -86,6 +87,6 @@ include("objectives.jl")
 
 function rand_device end
 
-include("sample.jl")
+include("rand_device.jl")
 
 end

src/rand_device.jl

@@ -0,0 +1,22 @@
+function rand_device(
+    rng::Random.AbstractRNG,
+    s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
+)
+    return Random.rand(rng, s)
+end
+
+function rand_device(
+    rng::Random.AbstractRNG,
+    s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
+    n::Int,
+)
+    return Random.rand(rng, s, n)
+end
+
+function rand_device(rng::Random.AbstractRNG, td::Bijectors.TransformedDistribution)
+    return Random.rand(rng, td)
+end
+
+function rand_device(rng::Random.AbstractRNG, td::Bijectors.TransformedDistribution, n::Int)
+    return Random.rand(rng, td, n)
+end