Merge pull request #2369 from Red-Portal/master

Add back missing `rng_from_array(::CuArray)` method
FluxML · Jan 24, 2024 · c2094d8 · c2094d8
2 parents 590adf2 + 80cd5fa
commit c2094d8
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 1 deletion.
diff --git a/ext/FluxCUDAExt/FluxCUDAExt.jl b/ext/FluxCUDAExt/FluxCUDAExt.jl
@@ -43,6 +43,7 @@ end
 ChainRulesCore.@non_differentiable check_use_cuda()
 
 include("functor.jl")
+include("utils.jl")
 
 function __init__()
     Flux.CUDA_LOADED[] = true

diff --git a/ext/FluxCUDAExt/utils.jl b/ext/FluxCUDAExt/utils.jl
@@ -1 +1 @@
-rng_from_array(::CuArray) = CUDA.default_rng()
+Flux.rng_from_array(::CuArray) = CUDA.default_rng()
diff --git a/test/ext_cuda/runtests.jl b/test/ext_cuda/runtests.jl
@@ -32,3 +32,7 @@ end
 @testset "ctc" begin
   include("ctc.jl")
 end
+@testset "utils" begin
+  include("utils.jl")
+end
+
diff --git a/test/ext_cuda/utils.jl b/test/ext_cuda/utils.jl
@@ -9,3 +9,9 @@
     @test y isa Wrapped{<:CuArray}
   end
 end
+
+@testset "rng_from_array" begin
+    x = cu(randn(2,2))
+    rng = Flux.rng_from_array(x)
+    @test rng == CUDA.default_rng()
+end
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		rng_from_array(::CuArray) = CUDA.default_rng()
		Flux.rng_from_array(::CuArray) = CUDA.default_rng()