// RUN: mlir-opt %s -shard-simplify | FileCheck %s

shard.grid @grid_ag(shape = 2x2)
shard.grid @grid_ag_alt(shape = 2x2)

// CHECK-LABEL: func.func @all_gather_all_slice_identity
func.func @all_gather_all_slice_identity(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
    : tensor<4x4xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 1
    : tensor<4x8xf32> -> tensor<4x4xf32>
  // CHECK-NOT: shard.all_gather
  // CHECK-NOT: shard.all_slice
  // CHECK: return %arg0 : tensor<4x4xf32>
  return %1 : tensor<4x4xf32>
}

// CHECK-LABEL: func.func @all_gather_all_slice_different_axis
func.func @all_gather_all_slice_different_axis(
    %arg0: tensor<4x4xf32>) -> tensor<2x8xf32> {
  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
    : tensor<4x4xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 0
    : tensor<4x8xf32> -> tensor<2x8xf32>
  // CHECK: shard.all_gather
  // CHECK: shard.all_slice
  return %1 : tensor<2x8xf32>
}

// CHECK-LABEL: func.func @all_gather_all_slice_different_grid_axes
func.func @all_gather_all_slice_different_grid_axes(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [0] gather_axis = 0
    : tensor<4x4xf32> -> tensor<8x4xf32>
  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 0
    : tensor<8x4xf32> -> tensor<4x4xf32>
  // CHECK: shard.all_gather
  // CHECK: shard.all_slice
  return %1 : tensor<4x4xf32>
}

// CHECK-LABEL: func.func @all_gather_all_slice_different_grid
func.func @all_gather_all_slice_different_grid(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
    : tensor<4x4xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid_ag_alt grid_axes = [1] slice_axis = 1
    : tensor<4x8xf32> -> tensor<4x4xf32>
  // CHECK: shard.all_gather
  // CHECK: shard.all_slice
  return %1 : tensor<4x4xf32>
}

// CHECK-LABEL: func.func @all_slice_all_gather_identity
func.func @all_slice_all_gather_identity(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
    : tensor<4x4xf32> -> tensor<4x2xf32>
  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 1
    : tensor<4x2xf32> -> tensor<4x4xf32>
  // CHECK-NOT: shard.all_slice
  // CHECK-NOT: shard.all_gather
  // CHECK: return %arg0 : tensor<4x4xf32>
  return %1 : tensor<4x4xf32>
}

// CHECK-LABEL: func.func @all_slice_all_gather_different_axis
func.func @all_slice_all_gather_different_axis(
    %arg0: tensor<4x4xf32>) -> tensor<8x2xf32> {
  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
    : tensor<4x4xf32> -> tensor<4x2xf32>
  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 0
    : tensor<4x2xf32> -> tensor<8x2xf32>
  // CHECK: shard.all_slice
  // CHECK: shard.all_gather
  return %1 : tensor<8x2xf32>
}

// CHECK-LABEL: func.func @all_slice_all_gather_different_grid
func.func @all_slice_all_gather_different_grid(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
    : tensor<4x4xf32> -> tensor<4x2xf32>
  %1 = shard.all_gather %0 on @grid_ag_alt grid_axes = [1] gather_axis = 1
    : tensor<4x2xf32> -> tensor<4x4xf32>
  // CHECK: shard.all_slice
  // CHECK: shard.all_gather
  return %1 : tensor<4x4xf32>
}

// CHECK-LABEL: func.func @all_slice_all_gather_different_grid_axes
func.func @all_slice_all_gather_different_grid_axes(
    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [0] slice_axis = 0
    : tensor<4x4xf32> -> tensor<2x4xf32>
  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 0
    : tensor<2x4xf32> -> tensor<4x4xf32>
  // CHECK: shard.all_slice
  // CHECK: shard.all_gather
  return %1 : tensor<4x4xf32>
}

shard.grid @grid0(shape = 4x2)
shard.grid @grid1(shape = 4)

// Checks that `all_reduce(x) + all_reduce(y)` gets transformed to
// `all_reduce(x + y)`.
// CHECK-LABEL: func.func @all_reduce_arith_addf_endomorphism
func.func @all_reduce_arith_addf_endomorphism(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ARG0]], %[[ARG1]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: %[[ALL_REDUCE_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ADD_RES]]
  // CHECK: return %[[ALL_REDUCE_RES]]
  return %2 : tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_addf_endomorphism_multiple_uses_of_result
func.func @all_reduce_arith_addf_endomorphism_multiple_uses_of_result(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ARG0]], %[[ARG1]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: %[[ALL_REDUCE_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ADD_RES]]
  // CHECK: return %[[ALL_REDUCE_RES]], %[[ALL_REDUCE_RES]]
  return %2, %2 : tensor<5xf32>, tensor<5xf32>
}

// Do not simplify if there is another use of one of the all-reduces.
// CHECK-LABEL: func.func @all_reduce_arith_addf_endomorphism_multiple_uses_of_all_reduce_result
func.func @all_reduce_arith_addf_endomorphism_multiple_uses_of_all_reduce_result(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
  // CHECK: %[[ALL_REDUCE_0_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG0]]
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ALL_REDUCE_1_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG1]]
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ALL_REDUCE_0_RES]], %[[ALL_REDUCE_1_RES]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: return %[[ALL_REDUCE_0_RES]], %[[ADD_RES]]
  return %0, %2 : tensor<5xf32>, tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_addf_no_endomorphism_different_grid
func.func @all_reduce_arith_addf_no_endomorphism_different_grid(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf32> {
  // CHECK: %[[ALL_REDUCE0:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG0]] on @grid0
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ALL_REDUCE1:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG1]] on @grid1
  %1 = shard.all_reduce %arg1 on @grid1 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ALL_REDUCE0]], %[[ALL_REDUCE1]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: return %[[ADD_RES]]
  return %2 : tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_addf_no_endomorphism_different_grid_axes
func.func @all_reduce_arith_addf_no_endomorphism_different_grid_axes(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf32> {
  // CHECK: %[[ALL_REDUCE0:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG0]] on @grid0 grid_axes = [0]
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ALL_REDUCE1:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG1]] on @grid0 grid_axes = [1]
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [1]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ALL_REDUCE0]], %[[ALL_REDUCE1]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: return %[[ADD_RES]]
  return %2 : tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_addf_no_endomorphism_wrong_reduction_kind
func.func @all_reduce_arith_addf_no_endomorphism_wrong_reduction_kind(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf32> {
  // CHECK: %[[ALL_REDUCE0:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG0]] on @grid0 grid_axes = [0] reduction = max
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] reduction = max
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ALL_REDUCE1:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG1]] on @grid0 grid_axes = [0]
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ALL_REDUCE0]], %[[ALL_REDUCE1]]
  %2 = arith.addf %0, %1 : tensor<5xf32>
  // CHECK: return %[[ADD_RES]]
  return %2 : tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_addf_no_endomorphism_different_operand_result_element_types
func.func @all_reduce_arith_addf_no_endomorphism_different_operand_result_element_types(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf64> {
  // CHECK: %[[ALL_REDUCE0:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG0]] on @grid0 grid_axes = [0]
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf64>
  // CHECK: %[[ALL_REDUCE1:[A-Za-z0-9_]*]] = shard.all_reduce %[[ARG1]] on @grid0 grid_axes = [0]
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0]
    : tensor<5xf32> -> tensor<5xf64>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.addf %[[ALL_REDUCE0]], %[[ALL_REDUCE1]]
  %2 = arith.addf %0, %1 : tensor<5xf64>
  // CHECK: return %[[ADD_RES]]
  return %2 : tensor<5xf64>
}

// Checks that `min(all_reduce(x), all_reduce(y))` gets transformed to
// `all_reduce(min(x, y))`.
// CHECK-LABEL: func.func @all_reduce_arith_minimumf_endomorphism
func.func @all_reduce_arith_minimumf_endomorphism(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg0: tensor<5xf32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xf32>
    %arg1: tensor<5xf32>) -> tensor<5xf32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] reduction = min
    : tensor<5xf32> -> tensor<5xf32>
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0] reduction = min
    : tensor<5xf32> -> tensor<5xf32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.minimumf %[[ARG0]], %[[ARG1]]
  %2 = arith.minimumf %0, %1 : tensor<5xf32>
  // CHECK: %[[ALL_REDUCE_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ADD_RES]] on @grid0 grid_axes = [0] reduction = min
  // CHECK: return %[[ALL_REDUCE_RES]]
  return %2 : tensor<5xf32>
}

// CHECK-LABEL: func.func @all_reduce_arith_minsi_endomorphism
func.func @all_reduce_arith_minsi_endomorphism(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<5xi32>
    %arg0: tensor<5xi32>,
    // CHECK-SAME: %[[ARG1:[A-Za-z0-9_]*]]: tensor<5xi32>
    %arg1: tensor<5xi32>) -> tensor<5xi32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] reduction = min
    : tensor<5xi32> -> tensor<5xi32>
  %1 = shard.all_reduce %arg1 on @grid0 grid_axes = [0] reduction = min
    : tensor<5xi32> -> tensor<5xi32>
  // CHECK: %[[ADD_RES:[A-Za-z0-9_]*]] = arith.minsi %[[ARG0]], %[[ARG1]]
  %2 = arith.minsi %0, %1 : tensor<5xi32>
  // CHECK: %[[ALL_REDUCE_RES:[A-Za-z0-9_]*]] = shard.all_reduce %[[ADD_RES]] on @grid0 grid_axes = [0] reduction = min
  // CHECK: return %[[ALL_REDUCE_RES]]
  return %2 : tensor<5xi32>
}

// Ensure this case without endomorphism op not crash.
// CHECK-LABEL: func.func @no_endomorphism_op
func.func @no_endomorphism_op(%arg0: tensor<2xi64>) -> i64 {
  %c0 = arith.constant 0 : index
  %c1_i64 = arith.constant 1 : i64
  // CHECK: tensor.extract
  %extracted = tensor.extract %arg0[%c0] : tensor<2xi64>
  // CHECK: arith.maxsi
  %0 = arith.maxsi %extracted, %c1_i64 : i64
  return %0 : i64
}

// -----
// AllReduceOp + AllSliceOp -> ReduceScatterOp tests
// -----

// Basic case: all_slice(all_reduce(x)) with matching grid and axes folds
// into reduce_scatter.
// CHECK-LABEL: func.func @all_reduce_all_slice_to_reduce_scatter
func.func @all_reduce_all_slice_to_reduce_scatter(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<1x8xf32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] : tensor<4x8xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid0 grid_axes = [0] slice_axis = 0 : tensor<4x8xf32> -> tensor<1x8xf32>
  // CHECK: %[[RS:.*]] = shard.reduce_scatter %[[ARG0]] on @grid0 grid_axes = [0] scatter_dim = 0
  // CHECK-SAME: : tensor<4x8xf32> -> tensor<1x8xf32>
  // CHECK: return %[[RS]]
  return %1 : tensor<1x8xf32>
}

// Verify non-default reduction kind is preserved.
// CHECK-LABEL: func.func @all_reduce_all_slice_max_reduction
func.func @all_reduce_all_slice_max_reduction(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<1x8xf32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] reduction = max : tensor<4x8xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid0 grid_axes = [0] slice_axis = 0 : tensor<4x8xf32> -> tensor<1x8xf32>
  // CHECK: %[[RS:.*]] = shard.reduce_scatter %[[ARG0]] on @grid0 grid_axes = [0] reduction = max scatter_dim = 0
  // CHECK-SAME: : tensor<4x8xf32> -> tensor<1x8xf32>
  // CHECK: return %[[RS]]
  return %1 : tensor<1x8xf32>
}

// Slice on a different tensor axis than the reduce axes.
// CHECK-LABEL: func.func @all_reduce_all_slice_different_slice_axis
func.func @all_reduce_all_slice_different_slice_axis(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<4x4xf32> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [1] : tensor<4x8xf32> -> tensor<4x8xf32>
  %1 = shard.all_slice %0 on @grid0 grid_axes = [1] slice_axis = 1 : tensor<4x8xf32> -> tensor<4x4xf32>
  // CHECK: %[[RS:.*]] = shard.reduce_scatter %[[ARG0]] on @grid0 grid_axes = [1] scatter_dim = 1
  // CHECK-SAME: : tensor<4x8xf32> -> tensor<4x4xf32>
  // CHECK: return %[[RS]]
  return %1 : tensor<4x4xf32>
}

// Do not fold when grids differ.
// CHECK-LABEL: func.func @all_reduce_all_slice_different_grid
func.func @all_reduce_all_slice_different_grid(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<1x8xf32> {
  // CHECK: %[[AR:.*]] = shard.all_reduce %[[ARG0]] on @grid0
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] : tensor<4x8xf32> -> tensor<4x8xf32>
  // CHECK: %[[AS:.*]] = shard.all_slice %[[AR]] on @grid1
  %1 = shard.all_slice %0 on @grid1 grid_axes = [0] slice_axis = 0 : tensor<4x8xf32> -> tensor<1x8xf32>
  // CHECK: return %[[AS]]
  return %1 : tensor<1x8xf32>
}

// Do not fold when grid_axes differ.
// CHECK-LABEL: func.func @all_reduce_all_slice_different_grid_axes
func.func @all_reduce_all_slice_different_grid_axes(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<4x4xf32> {
  // CHECK: %[[AR:.*]] = shard.all_reduce %[[ARG0]] on @grid0 grid_axes = [0]
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] : tensor<4x8xf32> -> tensor<4x8xf32>
  // CHECK: %[[AS:.*]] = shard.all_slice %[[AR]] on @grid0 grid_axes = [1]
  %1 = shard.all_slice %0 on @grid0 grid_axes = [1] slice_axis = 1 : tensor<4x8xf32> -> tensor<4x4xf32>
  // CHECK: return %[[AS]]
  return %1 : tensor<4x4xf32>
}

// Verify element type conversion is preserved (all_reduce input/output types may differ).
// CHECK-LABEL: func.func @all_reduce_all_slice_type_promotion
func.func @all_reduce_all_slice_type_promotion(
    // CHECK-SAME: %[[ARG0:[A-Za-z0-9_]*]]: tensor<4x8xf32>
    %arg0: tensor<4x8xf32>) -> tensor<1x8xf64> {
  %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0] : tensor<4x8xf32> -> tensor<4x8xf64>
  %1 = shard.all_slice %0 on @grid0 grid_axes = [0] slice_axis = 0 : tensor<4x8xf64> -> tensor<1x8xf64>
  // CHECK: %[[RS:.*]] = shard.reduce_scatter %[[ARG0]] on @grid0 grid_axes = [0] scatter_dim = 0
  // CHECK-SAME: : tensor<4x8xf32> -> tensor<1x8xf64>
  // CHECK: return %[[RS]]
  return %1 : tensor<1x8xf64>
}