Fixed similar and broadcasting wrt. supported_number_types(DeviceType)

Author: stemann

src/device.jl

@@ -3,6 +3,31 @@
     DeviceTypeGPU = 1
 end
 
+function supported_number_types(device_type::DeviceType = DeviceTypeCPU)
+    types = [
+        Bool,
+        UInt8,
+        UInt16,
+        UInt32,
+        UInt64,
+        Int8,
+        Int16,
+        Int32,
+        Int64,
+        # TODO Float16, 
+        Float32,
+        # TODO Core.BFloat16
+        ComplexF32,
+    ]
+    if device_type == DeviceTypeCPU
+        return vcat(types, [Float64])
+    elseif device_type == DeviceTypeGPU
+        return types
+    else
+        throw(ArgumentError("Unsupported device type: $device_type"))
+    end
+end
+
 mutable struct Device
     mlx_device::Wrapper.mlx_device
 

test/array_tests.jl

@@ -1,27 +1,25 @@
+@static if VERSION < v"1.11"
+    using ScopedValues
+else
+    using Base.ScopedValues
+end
+
 using MLX
 using Random
 using Test
 
 @testset "MLXArray" begin
     Random.seed!(42)
 
+    device_types = [MLX.DeviceTypeCPU]
+    if MLX.metal_is_available()
+        push!(device_types, MLX.DeviceTypeGPU)
+    end
+
+    element_types = MLX.supported_number_types()
+
     @test IndexStyle(MLXArray) == IndexLinear()
-    element_types = (
-        Bool,
-        UInt8,
-        UInt16,
-        UInt32,
-        UInt64,
-        Int8,
-        Int16,
-        Int32,
-        Int64,
-        # TODO Float16, 
-        Float32,
-        Float64,
-        # TODO Core.BFloat16
-        ComplexF32,
-    )
+
     array_sizes = [(), (1,), (2,), (1, 1), (2, 1), (2, 2), (1, 1, 1)]
     @testset "AbstractArray interface" begin
         for T in element_types, array_size in array_sizes
@@ -51,29 +49,44 @@ using Test
                 @test unsafe_wrap(mlx_array) == array
 
                 @testset "similar(::$MLXArray{$T, $N}), array_size=$array_size" begin
-                    similar_mlx_array = similar(mlx_array)
-                    @test typeof(similar_mlx_array) == typeof(mlx_array)
-                    @test size(similar_mlx_array) == size(mlx_array)
-                    @test similar_mlx_array !== mlx_array
+                    for device_type in device_types
+                        if T ∉ MLX.supported_number_types(device_type)
+                            continue
+                        end
+                        @testset "similar(::$MLXArray{$T, $N}), with array_size=$array_size, $device_type" begin
+                            with(MLX.device => MLX.Device(; device_type)) do
+                                similar_mlx_array = similar(mlx_array)
+                                @test typeof(similar_mlx_array) == typeof(mlx_array)
+                                @test size(similar_mlx_array) == size(mlx_array)
+                                @test similar_mlx_array !== mlx_array
+                            end
+                        end
+                    end
                 end
             end
         end
     end
+
     @testset "Unsupported Number types" begin
         @test_throws ArgumentError convert(MLX.Wrapper.mlx_dtype, Rational{Int})
     end
 
     @testset "Broadcasting interface" begin
-        for T in element_types, array_size in array_sizes
+        for device_type in device_types,
+            T in MLX.supported_number_types(device_type),
+            array_size in array_sizes
+
             N = length(array_size)
-            @testset "broadcast(identity, ::$MLXArray{$T, $N}), array_size=$array_size" begin
+            @testset "broadcast(identity, ::$MLXArray{$T, $N}), array_size=$array_size, $device_type" begin
                 array = rand(T, array_size)
                 mlx_array = MLXArray(array)
 
-                result = identity.(mlx_array)
-                @test result isa MLXArray
-                @test result == mlx_array
-                @test result !== mlx_array
+                with(MLX.device => MLX.Device(; device_type)) do
+                    result = identity.(mlx_array)
+                    @test result isa MLXArray
+                    @test result == mlx_array
+                    @test result !== mlx_array
+                end
             end
         end
     end

test/device_tests.jl

@@ -3,6 +3,7 @@
 else
     using Base.ScopedValues
 end
+
 using MLX
 using Test
 

test/stream_tests.jl

@@ -3,6 +3,7 @@
 else
     using Base.ScopedValues
 end
+
 using MLX
 using Test