Sharded return fix

src/Tracing.jl

@@ -1520,7 +1520,10 @@ Base.@nospecializeinfer function make_tracer(
                     (XLA.PJRT.AsyncEmptyBuffer,), size(prev), Sharding.NoShardInfo()
                 )
             else
-                error("TODO(#2230): implement sharding")
+                ndev = Sharding.ndevices(sharding)
+                res = ConcretePJRTArray{T,N,ndev}(
+                    ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev), size(prev), Sharding.NoShardInfo()
+                )
             end
             seen[prev] = res
             return res
@@ -1531,7 +1534,9 @@ Base.@nospecializeinfer function make_tracer(
                     XLA.IFRT.AsyncEmptyArray, size(prev), Sharding.NoShardInfo()
                 )
             else
-                error("TODO(#2230): implement sharding")
+                res = ConcreteIFRTArray{T,N}(
+                    XLA.IFRT.AsyncEmptyArray, size(prev), Sharding.NoShardInfo()
+                )
             end
             seen[prev] = res
             return res

src/compiler/Compiler.jl

@@ -1198,8 +1198,14 @@ function compile_mlir!(
         end
     end
 
+    sharding = if mlir_fn_res.is_sharded && !isempty(result_shardings_after_masking)
+        first(result_shardings_after_masking)
+    else
+        Sharding.NoSharding()
+    end
+
     concrete_result = make_tracer(
-        OrderedIdDict(), traced_result, ("result",), TracedToConcrete; runtime
+        OrderedIdDict(), traced_result, ("result",), TracedToConcrete; runtime, sharding
     )
 
     return Reactant.TracedUtils.CompiledMlirFnResult(

test/core/sharding/sharding.jl

@@ -663,3 +663,24 @@ end
     end
 end
 =#
+
+struct WrapTest{A}
+    data::A
+end
+
+wrap_f(a) = WrapTest(a .* a)
+
+@testset "Wrapped Sharded Array" begin
+    if length(addressable_devices) ≥ 2
+        ndev = length(addressable_devices)
+        mesh = Sharding.Mesh(reshape(addressable_devices, ndev), (:dev,))
+        sh = Sharding.DimsSharding(mesh, (1,), (:dev,))
+
+        x = Reactant.to_rarray(collect(Float64, 1:64); sharding = sh)
+        r = @jit(wrap_f(x))
+        @test r isa WrapTest
+        @test Array(r.data) ≈ collect(Float64, 1:64) .^ 2
+    else
+        @warn "Not enough addressable devices to run wrapped sharded array tests"
+    end
+end