From f455bb3ff57fc575b4d0d8629e44ce7e27ab513a Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 26 Mar 2025 10:11:47 +0100 Subject: [PATCH 1/5] Implemented broadcasting for MLXArray Also: * Added necessary eval of MLX array data in Base.unsafe_convert(::Type{Ptr{T}}, array::MLXArray{T, N}) * Implemented Base.similar for MLXArray. * Added supported_number_types(::DeviceType = DeviceTypeCPU) --- Project.toml | 3 ++- src/array.jl | 25 ++++++++++++++++++++++ test/array_tests.jl | 51 ++++++++++++++++++++++++++++++++++++++++++++ test/device_tests.jl | 1 + test/stream_tests.jl | 1 + 5 files changed, 80 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 507b6b0..4b0e56f 100644 --- a/Project.toml +++ b/Project.toml @@ -15,7 +15,8 @@ ScopedValues = "1" julia = "1" [extras] +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Random", "Test"] diff --git a/src/array.jl b/src/array.jl index 0c26f70..bb4deef 100644 --- a/src/array.jl +++ b/src/array.jl @@ -89,6 +89,13 @@ function Base.setindex!(array::MLXArray{T, N}, v::T, i::Int) where {T, N} return array end +function Base.similar(array::MLXArray{T, N}, ::Type{T}, ::Dims{N}) where {T, N} + stream = get_stream() + result_ref = Ref(Wrapper.mlx_array_new()) + Wrapper.mlx_zeros_like(result_ref, array.mlx_array, stream.mlx_stream) + return MLXArray{T, N}(result_ref[]) +end + # Strided array interface, cf. https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-strided-arrays function Base.strides(array::MLXArray) @@ -164,3 +171,21 @@ function Base.unsafe_wrap(array::MLXArray{T, N}) where {T, N} return PermutedDimsArray(wrapped_array, reverse(1:ndims(array))) end end + +# Broadcasting interface, cf. https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting + +Base.BroadcastStyle(::Type{<:MLXArray}) = Broadcast.ArrayStyle{MLXArray}() + +function Base.similar( + bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MLXArray}}, ::Type{TElement} +) where {TElement} + first_mlx_array(bc::Broadcast.Broadcasted) = first_mlx_array(bc.args) + function first_mlx_array(args::Tuple) + return first_mlx_array(first_mlx_array(args[1]), Base.tail(args)) + end + first_mlx_array(x) = x + first_mlx_array(::Tuple{}) = nothing + first_mlx_array(a::MLXArray, _) = a + first_mlx_array(::Any, rest) = first_mlx_array(rest) + return similar(first_mlx_array(bc)) +end diff --git a/test/array_tests.jl b/test/array_tests.jl index 025688a..4b2843b 100644 --- a/test/array_tests.jl +++ b/test/array_tests.jl @@ -1,7 +1,21 @@ +@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 + @test IndexStyle(MLXArray) == IndexLinear() array_sizes = [(), (1,), (2,), (1, 1), (2, 1), (3, 2), (4, 3, 2)] @@ -31,6 +45,22 @@ using Test array[1] = T(1) @test setindex!(mlx_array, T(1), 1) == array end + + @testset "similar(::$MLXArray{$T, $N}), array_size=$array_size" begin + 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 @@ -62,7 +92,28 @@ using Test @test Base.elsize(MLXArray{T, 0}) == Base.elsize(Array{T, 0}) end end + @testset "Unsupported Number types" begin @test_throws ArgumentError convert(MLX.Wrapper.mlx_dtype, Rational{Int}) end + + @testset "Broadcasting interface" begin + 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, $device_type" begin + array = rand(T, array_size) + mlx_array = MLXArray(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 end diff --git a/test/device_tests.jl b/test/device_tests.jl index 4051cab..3a7dbab 100644 --- a/test/device_tests.jl +++ b/test/device_tests.jl @@ -3,6 +3,7 @@ else using Base.ScopedValues end + using MLX using Test diff --git a/test/stream_tests.jl b/test/stream_tests.jl index 28b106f..b5a1429 100644 --- a/test/stream_tests.jl +++ b/test/stream_tests.jl @@ -3,6 +3,7 @@ else using Base.ScopedValues end + using MLX using Test From 428e6d353b43589c85811a46eacfed779015d9ea Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 3 Dec 2025 20:01:03 +0100 Subject: [PATCH 2/5] Added binary op broadcast tests * Added handling of broadcast over empty tuple. * Added constructors for UndefInitializer. * Added handling of https://github.com/JuliaLang/julia/issues/28866 - dropping 0-dim arrays to scalars. --- src/array.jl | 31 ++++++++++++++++++++++++++++++- test/array_tests.jl | 30 ++++++++++++++++++++++++------ 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/src/array.jl b/src/array.jl index bb4deef..716432d 100644 --- a/src/array.jl +++ b/src/array.jl @@ -66,6 +66,21 @@ MLXMatrix(array::AbstractMatrix{T}) where {T} = MLXMatrix{T}(array) const MLXVecOrMat{T} = Union{MLXVector{T}, MLXMatrix{T}} +# UndefInitializer + +function MLXArray{T, N}(::UndefInitializer, dims::Dims{N}) where {T, N} + stream = get_stream() + result_ref = Ref(Wrapper.mlx_array_new()) + shape = collect(Cint.(dims)) + dtype = convert(Wrapper.mlx_dtype, T) + Wrapper.mlx_zeros(result_ref, pointer(shape), Cint(N), dtype, stream.mlx_stream) + return MLXArray{T, N}(result_ref[]) +end + +function MLXArray{T}(::UndefInitializer, dims::Dims{N}) where {T, N} + return MLXArray{T, N}(undef, dims) +end + # AbstractArray interface, cf. https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array function Base.size(array::MLXArray) @@ -187,5 +202,19 @@ function Base.similar( first_mlx_array(::Tuple{}) = nothing first_mlx_array(a::MLXArray, _) = a first_mlx_array(::Any, rest) = first_mlx_array(rest) - return similar(first_mlx_array(bc)) + mlx_array = first_mlx_array(bc) + if isnothing(mlx_array) + return similar(MLXArray{TElement}, ()) + end + return similar(mlx_array) +end + +function Base.Broadcast.materialize( + bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MLXArray}} +) + result = copy(Broadcast.instantiate(bc)) + if iszero(ndims(result)) # Drop 0-dim arrays to scalars, cf. https://github.com/JuliaLang/julia/issues/28866 + return result[] + end + return result end diff --git a/test/array_tests.jl b/test/array_tests.jl index 4b2843b..de299bc 100644 --- a/test/array_tests.jl +++ b/test/array_tests.jl @@ -98,20 +98,38 @@ using Test end @testset "Broadcasting interface" begin + @testset "broadcast over tuple with no MLXArray" begin + result = similar( + Broadcast.Broadcasted{Broadcast.ArrayStyle{MLXArray}}(identity, ()), Bool + ) + @test result isa MLXArray{Bool, 0} + end + + test_cases(T, array_size) = [ + (fn = identity, args = ()), + (fn = T == Bool ? xor : +, args = (T == Bool ? true : T(2),)), + (fn = T == Bool ? xor : +, args = (MLXArray(rand(T, array_size)),)), + ] for device_type in device_types, T in MLX.supported_number_types(device_type), - array_size in array_sizes + array_size in array_sizes, + test_case in test_cases(T, array_size) + compare_op = T <: Integer ? (==) : ≈ N = length(array_size) - @testset "broadcast(identity, ::$MLXArray{$T, $N}), array_size=$array_size, $device_type" begin + @testset "broadcast($(repr(test_case.fn)), $(join(map(arg -> "::$(typeof(arg))", [MLXArray{T, N}, test_case.args...]), ", "))), array_size=$array_size, $device_type" begin array = rand(T, array_size) mlx_array = MLXArray(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 + actual = broadcast(test_case.fn, mlx_array, test_case.args...) + expected = broadcast(test_case.fn, array, test_case.args...) + if N == 0 + @test actual isa T + else + @test actual isa MLXArray{T, N} + end + @test compare_op(actual, expected) end end end From f6a887bc29fc3166a0d506664c6fcb975d0a5213 Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Mon, 25 Aug 2025 08:48:54 +0200 Subject: [PATCH 3/5] Added constructors for BitArray et al. --- src/array.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/array.jl b/src/array.jl index 716432d..210271e 100644 --- a/src/array.jl +++ b/src/array.jl @@ -81,6 +81,14 @@ function MLXArray{T}(::UndefInitializer, dims::Dims{N}) where {T, N} return MLXArray{T, N}(undef, dims) end +# BitArray + +MLXArray(array::BitArray{N}) where {N} = MLXArray(Array{Bool}(array)) + +MLXVector(array::BitVector) = MLXVector(Vector{Bool}(array)) + +MLXMatrix(array::BitMatrix) = MLXMatrix(Matrix{Bool}(array)) + # AbstractArray interface, cf. https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array function Base.size(array::MLXArray) From 770e6fbd52ae682b3a4dac9689b1716ab498a31e Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 3 Dec 2025 20:28:02 +0100 Subject: [PATCH 4/5] Added BitArray tests --- test/array_tests.jl | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/array_tests.jl b/test/array_tests.jl index de299bc..33bdd52 100644 --- a/test/array_tests.jl +++ b/test/array_tests.jl @@ -97,6 +97,23 @@ using Test @test_throws ArgumentError convert(MLX.Wrapper.mlx_dtype, Rational{Int}) end + @testset "BitArray" begin + for array_size in array_sizes + N = length(array_size) + @testset "$MLXArray{Bool, $N}(::BitArray), array_size=$array_size" begin + array = BitArray(rand(Bool, array_size)) + if N > 2 || N == 0 + mlx_array = MLXArray(array) + elseif N > 1 + mlx_array = MLXMatrix(array) + else + mlx_array = MLXVector(array) + end + @test array == mlx_array + end + end + end + @testset "Broadcasting interface" begin @testset "broadcast over tuple with no MLXArray" begin result = similar( From a81b295ddfa44613d774a97edf4f67c0849d75da Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 3 Dec 2025 20:44:41 +0100 Subject: [PATCH 5/5] Ensure explicit type constructors are defined as well for BitArray --- src/array.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/array.jl b/src/array.jl index 210271e..5672a66 100644 --- a/src/array.jl +++ b/src/array.jl @@ -83,11 +83,9 @@ end # BitArray -MLXArray(array::BitArray{N}) where {N} = MLXArray(Array{Bool}(array)) +MLXArray{Bool, N}(array::BitArray{N}) where {N} = MLXArray(Array{Bool}(array)) -MLXVector(array::BitVector) = MLXVector(Vector{Bool}(array)) - -MLXMatrix(array::BitMatrix) = MLXMatrix(Matrix{Bool}(array)) +MLXArray(array::BitArray{N}) where {N} = MLXArray{Bool, N}(array) # AbstractArray interface, cf. https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array