test: functions and closures

LuxDL · Oct 18, 2024 · 37429fb · 37429fb
1 parent b3bef22
commit 37429fb
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 0 deletions.
diff --git a/test/amdgpu_tests.jl b/test/amdgpu_tests.jl
@@ -126,6 +126,27 @@ using FillArrays, Zygote  # Extensions
     end
 end
 
+@testset "Functions" begin
+    if MLDataDevices.functional(AMDGPUDevice)
+        @test get_device(tanh) isa MLDataDevices.UnknownDevice
+        @test get_device_type(tanh) <: MLDataDevices.UnknownDevice
+
+        f(x, y) = () -> (x, x .^ 2, y)
+
+        ff = f([1, 2, 3], 1)
+        @test get_device(ff) isa CPUDevice
+        @test get_device_type(ff) <: CPUDevice
+
+        ff_xpu = ff |> AMDGPUDevice()
+        @test get_device(ff_xpu) isa AMDGPUDevice
+        @test get_device_type(ff_xpu) <: AMDGPUDevice
+
+        ff_cpu = ff_xpu |> cpu_device()
+        @test get_device(ff_cpu) isa CPUDevice
+        @test get_device_type(ff_cpu) <: CPUDevice
+    end
+end
+
 @testset "Wrapped Arrays" begin
     if MLDataDevices.functional(AMDGPUDevice)
         x = rand(10, 10) |> AMDGPUDevice()

diff --git a/test/cuda_tests.jl b/test/cuda_tests.jl
@@ -151,6 +151,27 @@ using FillArrays, Zygote  # Extensions
     end
 end
 
+@testset "Functions" begin
+    if MLDataDevices.functional(CUDADevice)
+        @test get_device(tanh) isa MLDataDevices.UnknownDevice
+        @test get_device_type(tanh) <: MLDataDevices.UnknownDevice
+
+        f(x, y) = () -> (x, x .^ 2, y)
+
+        ff = f([1, 2, 3], 1)
+        @test get_device(ff) isa CPUDevice
+        @test get_device_type(ff) <: CPUDevice
+
+        ff_xpu = ff |> CUDADevice()
+        @test get_device(ff_xpu) isa CUDADevice
+        @test get_device_type(ff_xpu) <: CUDADevice
+
+        ff_cpu = ff_xpu |> cpu_device()
+        @test get_device(ff_cpu) isa CPUDevice
+        @test get_device_type(ff_cpu) <: CPUDevice
+    end
+end
+
 @testset "Wrapped Arrays" begin
     if MLDataDevices.functional(CUDADevice)
         x = rand(10, 10) |> CUDADevice()

diff --git a/test/metal_tests.jl b/test/metal_tests.jl
@@ -115,6 +115,27 @@ using FillArrays, Zygote  # Extensions
     end
 end
 
+@testset "Functions" begin
+    if MLDataDevices.functional(MetalDevice)
+        @test get_device(tanh) isa MLDataDevices.UnknownDevice
+        @test get_device_type(tanh) <: MLDataDevices.UnknownDevice
+
+        f(x, y) = () -> (x, x .^ 2, y)
+
+        ff = f([1, 2, 3], 1)
+        @test get_device(ff) isa CPUDevice
+        @test get_device_type(ff) <: CPUDevice
+
+        ff_xpu = ff |> MetalDevice()
+        @test get_device(ff_xpu) isa MetalDevice
+        @test get_device_type(ff_xpu) <: MetalDevice
+
+        ff_cpu = ff_xpu |> cpu_device()
+        @test get_device(ff_cpu) isa CPUDevice
+        @test get_device_type(ff_cpu) <: CPUDevice
+    end
+end
+
 @testset "Wrapper Arrays" begin
     if MLDataDevices.functional(MetalDevice)
         x = rand(Float32, 10, 10) |> MetalDevice()

diff --git a/test/oneapi_tests.jl b/test/oneapi_tests.jl
@@ -115,6 +115,27 @@ using FillArrays, Zygote  # Extensions
     end
 end
 
+@testset "Functions" begin
+    if MLDataDevices.functional(oneAPIDevice)
+        @test get_device(tanh) isa MLDataDevices.UnknownDevice
+        @test get_device_type(tanh) <: MLDataDevices.UnknownDevice
+
+        f(x, y) = () -> (x, x .^ 2, y)
+
+        ff = f([1, 2, 3], 1)
+        @test get_device(ff) isa CPUDevice
+        @test get_device_type(ff) <: CPUDevice
+
+        ff_xpu = ff |> oneAPIDevice()
+        @test get_device(ff_xpu) isa oneAPIDevice
+        @test get_device_type(ff_xpu) <: oneAPIDevice
+
+        ff_cpu = ff_xpu |> cpu_device()
+        @test get_device(ff_cpu) isa CPUDevice
+        @test get_device_type(ff_cpu) <: CPUDevice
+    end
+end
+
 @testset "Wrapper Arrays" begin
     if MLDataDevices.functional(oneAPIDevice)
         x = rand(10, 10) |> oneAPIDevice()

diff --git a/test/xla_tests.jl b/test/xla_tests.jl
@@ -114,6 +114,27 @@ using FillArrays, Zygote  # Extensions
     end
 end
 
+@testset "Functions" begin
+    if MLDataDevices.functional(XLADevice)
+        @test get_device(tanh) isa MLDataDevices.UnknownDevice
+        @test get_device_type(tanh) <: MLDataDevices.UnknownDevice
+
+        f(x, y) = () -> (x, x .^ 2, y)
+
+        ff = f([1, 2, 3], 1)
+        @test get_device(ff) isa CPUDevice
+        @test get_device_type(ff) <: CPUDevice
+
+        ff_xpu = ff |> XLADevice()
+        @test get_device(ff_xpu) isa XLADevice
+        @test get_device_type(ff_xpu) <: XLADevice
+
+        ff_cpu = ff_xpu |> cpu_device()
+        @test get_device(ff_cpu) isa CPUDevice
+        @test get_device_type(ff_cpu) <: CPUDevice
+    end
+end
+
 @testset "Wrapped Arrays" begin
     if MLDataDevices.functional(XLADevice)
         x = rand(10, 10) |> XLADevice()