Added test for functor Cholesky.

FluxML · Apr 25, 2020 · d5cb27e · d5cb27e
1 parent 621b4cb
commit d5cb27e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/src/functor.jl b/src/functor.jl
@@ -1,4 +1,5 @@
 import Adapt: adapt, adapt_storage
+using  LinearAlgebra: Cholesky
 using Zygote: IdSet
 
 functor(x) = (), _ -> x
@@ -115,3 +116,6 @@ paramtype(T::Type{<:Real}, m) = fmap(x -> adapt(T, x), m)
 
 f32(m) = paramtype(Float32, m)
 f64(m) = paramtype(Float64, m)
+
+# Functors for certain Julia data structures
+@functor Cholesky
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -1,5 +1,3 @@
-using LinearAlgebra
-
 """
     Chain(layers...)
 
@@ -260,5 +258,3 @@ end
 function Base.show(io::IO, b::SkipConnection)
   print(io, "SkipConnection(", b.layers, ", ", b.connection, ")")
 end
-
-@functor Cholesky
diff --git a/test/cuda/cuda.jl b/test/cuda/cuda.jl
@@ -1,6 +1,7 @@
 using Flux, Test
 using Flux.CuArrays
-using Flux: gpu
+using Flux: cpu, gpu
+using LinearAlgebra: I, cholesky, Cholesky
 
 @info "Testing GPU Support"
 
@@ -65,6 +66,17 @@ end
   @test gradient(foo, cu(rand(1)))[1] isa CuArray
 end
 
+@testset "GPU functors" begin
+  @testset "Cholesky" begin
+    M = 2.0*I(10) |> collect
+    Q = cholesky(M)
+    Q_gpu = Q |> gpu
+    @test Q_gpu isa Cholesky{<:Any,<:CuArray}
+    Q_cpu = Q_gpu |> cpu
+    @test Q_cpu == cholesky(eltype(Q_gpu).(M))
+  end
+end
+
 if CuArrays.has_cudnn()
   @info "Testing Flux/CUDNN"
   include("cudnn.jl")