gpu fix (#12)

make dists work on gpu :)

Manifest.toml

@@ -87,9 +87,9 @@ version = "1.1.0"
 
 [[DataStructures]]
 deps = ["InteractiveUtils", "OrderedCollections"]
-git-tree-sha1 = "a1b652fb77ae8ca7ea328fa7ba5aa151036e5c10"
+git-tree-sha1 = "f784254f428fb8fd7ac15982e5862a38a44523d3"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.17.6"
+version = "0.17.7"
 
 [[Dates]]
 deps = ["Printf"]
@@ -180,9 +180,9 @@ version = "0.7.2"
 
 [[LLVM]]
 deps = ["CEnum", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "74fe444b8b6d1ac01d639b2f9eaf395bcc2e24fc"
+git-tree-sha1 = "1d08d7e4250f452f6cb20e4574daaebfdbee0ff7"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "1.3.2"
+version = "1.3.3"
 
 [[LibGit2]]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

Project.toml

@@ -4,6 +4,7 @@ authors = ["Niklas Heim <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

src/cmean_gaussian.jl

@@ -29,18 +29,21 @@ julia> rand(p, ones(2))
  0.0767166501426535
 ```
 """
-struct CMeanGaussian{V<:AbstractVar,S<:AbstractArray,M} <: AbstractCGaussian
+struct CMeanGaussian{V<:AbstractVar,M,S<:AbstractArray} <: AbstractCGaussian
     mapping::M
     σ::S
     xlength::Int
     _nograd::Dict{Symbol,Bool}
 end
 
+CMeanGaussian{V}(m::M, σ::S, xlength::Int, d::Dict{Symbol,Bool}) where {V,M,S} =
+    CMeanGaussian{V,M,S}(m,σ,xlength,d)
+
 function CMeanGaussian{V}(m::M, σ, xlength::Int) where {V,M}
     _nograd = Dict(:σ => σ isa NoGradArray)
     σ = _nograd[:σ] ? σ.data : σ
     S = typeof(σ)
-    CMeanGaussian{V,S,M}(m, σ, xlength, _nograd)
+    CMeanGaussian{V,M,S}(m, σ, xlength, _nograd)
 end
 
 CMeanGaussian{DiagVar}(m, σ) = CMeanGaussian{DiagVar}(m, σ, size(σ,1))
@@ -62,10 +65,10 @@ end
 mean_var(p::CMeanGaussian, z::AbstractArray) = (mean(p, z), variance(p, z))
 
 # make sure that parameteric constructor is called...
-function Flux.functor(p::CMeanGaussian{V,S,M}) where {V,S,M}
+function Flux.functor(p::CMeanGaussian{V,M,S}) where {V,M,S}
     fs = fieldnames(typeof(p))
     nt = (; (name=>getfield(p, name) for name in fs)...)
-    nt, y -> CMeanGaussian{V,S,M}(y...)
+    nt, y -> CMeanGaussian{V}(y...)
 end
 
 function Flux.trainable(p::CMeanGaussian)

src/cmeanvar_gaussian.jl

@@ -54,10 +54,10 @@ function mean_var(p::CMeanVarGaussian{ScalarVar}, z::AbstractArray)
 end
 
 # make sure that parameteric constructor is called...
-function Flux.functor(p::CMeanVarGaussian{V,M}) where {V,M}
+function Flux.functor(p::CMeanVarGaussian{V}) where V
     fs = fieldnames(typeof(p))
     nt = (; (name=>getfield(p, name) for name in fs)...)
-    nt, y -> CMeanVarGaussian{V,M}(y...)
+    nt, y -> CMeanVarGaussian{V}(y...)
 end
 
 function Base.show(io::IO, p::CMeanVarGaussian{V}) where V

test/runtests.jl

@@ -8,7 +8,7 @@ if Flux.use_cuda[] using CuArrays end
 
 include("abstract_pdf.jl")
 include("gaussian.jl")
-include("cmean_gaussian.jl")
+include("nogradarray.jl")
 include("cmeanvar_gaussian.jl")
+include("cmean_gaussian.jl")
 include("constspec_gaussian.jl")
-include("nogradarray.jl")