feat: support implicit padding from XLA

Author: avik-pal

src/Compiler.jl

@@ -1167,14 +1167,16 @@ function codegen_flatten!(
 
         if is_sharded
             carg = inv_seen_args[arg]
+            condensed_op_sharding = convert(
+                Reactant.Sharding.XLA.CondensedOpSharding, linear_parameter_shardings[i]
+            )
             if Reactant.Sharding.is_sharded(carg)
-                # Currently disabling the error since we roundtrip from MHLO to generate
-                # the shardings
-                # # Check if the sharding provided is same as the one we have
-                # arg_condensed_op_sharding = Reactant.Sharding.XLA.CondensedOpSharding(
-                #     Reactant.Sharding.ShardingWithShape(carg.sharding, size(carg))
-                # )
-                # @assert arg_condensed_op_sharding == condensed_op_sharding "Sharding provided by the user ($arg_condensed_op_sharding) does not match the sharding computed by XLA ($condensed_op_sharding). This generally means that Reactant.jl made an error in generating the executable. Please open an issue with the error message and an MWE."
+                arg_condensed_op_sharding = convert(
+                    Reactant.Sharding.XLA.CondensedOpSharding,
+                    carg.sharding.sharding.hlo_sharding,
+                )
+                # Check if the sharding provided is same as the one we have
+                @assert arg_condensed_op_sharding == condensed_op_sharding "Sharding provided by the user ($arg_condensed_op_sharding) does not match the sharding computed by XLA ($condensed_op_sharding). This generally means that Reactant.jl made an error in generating the executable. Please open an issue with the error message and an MWE."
 
                 push!(flatten_code, :($usbuf = $flatcode.data))
                 for j in 1:length(mesh)
@@ -1183,9 +1185,6 @@ function codegen_flatten!(
                     push!(flatten_code, :($sbuf = XLA.synced_buffer(getindex($usbuf, $j))))
                 end
             else
-                condensed_op_sharding = convert(
-                    Reactant.Sharding.XLA.CondensedOpSharding, linear_parameter_shardings[i]
-                )
                 push!(flatten_code, :($usbuf = $flatcode))
                 device_to_array_slices = XLA.sharding_to_concrete_array_indices(
                     condensed_op_sharding, size(carg), mesh.logical_device_ids

src/xla/Sharding.jl

@@ -257,25 +257,28 @@ function sharding_to_concrete_array_indices(
         @assert length(partitions) == length(shape)
         shape = reverse(shape)
 
+        # XLA will automatically pad the inputs that don't match the final shape
+        partitionable_shape = map(zip(shape, partitions)) do (dim, n_shards)
+            dim % n_shards == 0 && return dim
+            res = dim + n_shards ÷ 2
+            return res - res % n_shards
+        end
+        partitionable_shape = Tuple(partitionable_shape)
+
         # Calculate indices for each dimension
-        axis_indices = map(zip(shape, partitions)) do (dim, n_shards)
-            @assert dim > 0 "Invalid dimension: $dim"
-            @assert n_shards > 0 "Invalid number of shards: $n_shards"
-            n_shards == 1 && return [1:dim]
-            shard_size, remainder = divrem(dim, n_shards)
-
-            if remainder != 0
-                throw(
-                    DimensionMismatch(
-                        "Dimension of Size $(dim) cannot be partitioned into $(n_shards) \
-                         shards each of size $(shard_size) (remainder = $(remainder)).",
-                    ),
-                )
+        axis_indices =
+            map(zip(partitionable_shape, shape, partitions)) do (dim_padded, dim, n_shards)
+                @assert dim > 0 "Invalid dimension: $dim"
+                @assert n_shards > 0 "Invalid number of shards: $n_shards"
+                n_shards == 1 && return [1:dim]
+                shard_size = dim_padded ÷ n_shards
+
+                return [
+                    (i * shard_size + 1):min((i + 1) * shard_size, dim) for
+                    i in 0:(n_shards - 1)
+                ]
             end
 
-            return [(i * shard_size + 1):((i + 1) * shard_size) for i in 0:(n_shards - 1)]
-        end
-
         indices = Dict{Int,NTuple{N,UnitRange{Int}}}()
         device_idx = 1
         for _ in 1:num_replicas