fix: check for ifrt_array_copy_to_host_buffer

Author: avik-pal

src/xla/IFRT/Array.jl

@@ -138,7 +138,11 @@ function XLA.buffer_on_cpu(::Array)
 end
 
 function XLA.to_host(buffer::Array, data, reactant_sharding)
-    if length(XLA.devices(XLA.sharding(buffer))) == 1
+    reactant_sharding = Reactant.Sharding.unwrap_shardinfo(reactant_sharding)
+
+    if is_single_device_sharding(XLA.sharding(buffer)) ||
+        is_fully_replicated(XLA.sharding(buffer)) ||
+        reactant_sharding isa Reactant.Sharding.NoSharding
         GC.@preserve buffer data begin
             @ccall MLIR.API.mlir_c.ifrt_array_copy_to_host_buffer(
                 buffer.buffer::Ptr{Cvoid}, data::Ptr{Cvoid}
@@ -147,7 +151,6 @@ function XLA.to_host(buffer::Array, data, reactant_sharding)
         return data
     end
 
-    reactant_sharding = Reactant.Sharding.unwrap_shardinfo(reactant_sharding)
     @assert reactant_sharding isa Reactant.Sharding.HloSharding
     client = XLA.client(buffer)
     all_devices = XLA.get_device.((client,), reactant_sharding.mesh.device_ids)