Merge pull request #483 from ChrisRackauckas-Claude/add-promote-eltype-gpu

Add promote_eltype for CuArray / ROCArray / MtlArray

Author: ChrisRackauckas

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "7.23.0"
+version = "7.24.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/ArrayInterfaceAMDGPUExt.jl

@@ -12,4 +12,10 @@ end
 
 ArrayInterface.device(::Type{<:AMDGPU.ROCArray}) = ArrayInterface.GPU()
 
+function ArrayInterface.promote_eltype(
+        ::Type{<:AMDGPU.ROCArray{T, N, B}}, ::Type{T2}
+    ) where {T, N, B, T2}
+    return AMDGPU.ROCArray{promote_type(T, T2), N, B}
+end
+
 end # module

ext/ArrayInterfaceCUDAExt.jl

@@ -13,4 +13,10 @@ end
 
 ArrayInterface.device(::Type{<:CUDA.CuArray}) = ArrayInterface.GPU()
 
+function ArrayInterface.promote_eltype(
+        ::Type{<:CUDA.CuArray{T, N, M}}, ::Type{T2}
+    ) where {T, N, M, T2}
+    return CUDA.CuArray{promote_type(T, T2), N, M}
+end
+
 end # module

ext/ArrayInterfaceMetalExt.jl

@@ -12,4 +12,10 @@ end
 
 ArrayInterface.device(::Type{<:Metal.MtlArray}) = ArrayInterface.GPU()
 
+function ArrayInterface.promote_eltype(
+        ::Type{<:Metal.MtlArray{T, N, S}}, ::Type{T2}
+    ) where {T, N, S, T2}
+    return Metal.MtlArray{promote_type(T, T2), N, S}
+end
+
 end # module