fix: dispatches for GB

Author: avik-pal

src/ConcreteRArray.jl

@@ -393,14 +393,6 @@ buffer_on_cpu(::Any) = true
 buffer_on_cpu(x::ConcretePJRTArray) = all(XLA.buffer_on_cpu, x.data)
 buffer_on_cpu(x::ConcreteIFRTArray) = XLA.buffer_on_cpu(x.data)
 
-function Ops.constant(x::AbstractConcreteArray; kwargs...)
-    return Ops.constant(Base.convert(Array, x); kwargs...)
-end
-
-function Ops.constant(x::AbstractConcreteNumber{T}; kwargs...) where {T}
-    return Ops.constant(Base.convert(T, x); kwargs...)
-end
-
 function Base.zero(x::ConcretePJRTArray{T,N}) where {T,N}
     return ConcretePJRTArray(
         zeros(T, size(x)...); client=XLA.client(x), device=XLA.device(x), x.sharding
@@ -464,3 +456,9 @@ function Base.mapreducedim!(
     fn(f, op, R, A)
     return R
 end
+
+function Base.map!(f, R::Union{AnyConcreteIFRTArray,AnyConcretePJRTArray}, A::AbstractArray)
+    fn = compile(Base.map!, (f, R, A))
+    fn(f, R, A)
+    return R
+end

src/Ops.jl

@@ -122,6 +122,16 @@ end
     end
 end
 
+@noinline function constant(
+    x::AbstractArray{T,N}; location=mlir_stacktrace("constant", @__FILE__, @__LINE__)
+) where {T,N}
+    return constant(collect(x); location)
+end
+
+@noinline function constant(x::Reactant.AbstractConcreteArray; kwargs...)
+    return constant(Base.convert(Array, x); kwargs...)
+end
+
 @noinline function constant(
     x::T; location=mlir_stacktrace("constant", @__FILE__, @__LINE__)
 ) where {T<:Number}
@@ -130,6 +140,10 @@ end
     return TracedRNumber{T}((), res.mlir_data)
 end
 
+@noinline function constant(x::Reactant.AbstractConcreteNumber{T}; kwargs...) where {T}
+    return constant(Base.convert(T, x); kwargs...)
+end
+
 function fill(
     v, dims::Base.DimOrInd...; location=mlir_stacktrace("fill", @__FILE__, @__LINE__)
 )

src/TracedRArray.jl

@@ -1084,4 +1084,11 @@ function Base.findmax(f, x::AnyTracedRArray; dims::Union{Integer,Nothing}=nothin
     return (values, linear_indices)
 end
 
+Base.map(f, x::AnyTracedRArray) = f.(x)
+
+function Base.map!(f, y::AnyTracedRArray, x::AbstractArray)
+    y .= f.(x)
+    return y
+end
+
 end