tf_device.return %5 : tensor<512xi32>
    }) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<512xi32>, tensor<1024xi32>, tensor<1024xi32>, tensor<1024xf32>)
    // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:4 = "tf_device.launch"
    // CHECK: "tf._TPUCompileMlir"()
    // CHECK: is_bounded_dynamic_dim: true

  func.return %0 : tensor<1024xf32>

  // 1.0 * 2.0 * 512 + 4.0 = 1028.0
  // CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<1024xf32>
  // CHECK:  return %[[CST]]
}

// CHECK-LABEL: @ConstantFoldFullyConnectedNoBias
func.func @ConstantFoldFullyConnectedNoBias() -> tensor<1024xf32> {

// CHECK-SAME:                                      %[[VAL_0:.*]]: tensor<10x8xi32>,
// CHECK-SAME:                                      %[[VAL_1:.*]]: tensor<10x8xi32>) -> tensor<10x8xi32> {
// CHECK:           %[[VAL_2:.*]] = "tf.Mod"(%[[VAL_0]], %[[VAL_1]]) : (tensor<10x8xi32>, tensor<10x8xi32>) -> tensor<10x8xi32>
// CHECK:           return %[[VAL_2]] : tensor<10x8xi32>
// CHECK:         }

// CHECK-LABEL: @random_shuffle_1D_10240
func.func @random_shuffle_1D_10240(%input: tensor<10240xf32>) -> tensor<10240xf32> {
  // CHECK: mhlo.rng{{.*UNIFORM.*}}
  // CHECK: mhlo.sort
  // CHECK: mhlo.rng{{.*UNIFORM.*}}
  // CHECK: mhlo.sort
  %0 = "tf.RandomShuffle"(%input) : (tensor<10240xf32>) -> (tensor<10240xf32>)
  func.return %0: tensor<10240xf32>
}

// -----

// CHECK-LABEL: @random_shuffle_3D

tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>,...

  %cst1 = arith.constant dense<2.0> : tensor<1024xf32>
  %cst2 = arith.constant dense<2.0> : tensor<1024xf32>

  %0 = "tfl.fully_connected"(%arg0, %cst0, %cst1) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<128x256000xf32>, tensor<1024x256000xf32>, tensor<1024xf32>) -> tensor<128x1024xf32>

  // CHECK: "tfl.pack"(%arg0, %arg1) <{axis = -2 : i32, values_count = 2 : i32}>
  %0 = "tfl.pack"(%arg0, %arg1) {axis = -2 : i32, values_count = 2 : i32} : (tensor<1x4xi32>, tensor<1x4xi32>) -> tensor<1x2x4xi32>
  func.return %0 : tensor<1x2x4xi32>
}

func.func @packNegInputAxis3(%arg0: tensor<1x4xi32>, %arg1: tensor<1x4xi32>) -> tensor<2x1x4xi32> {

  // CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
  // CHECK: %[[SUBGRAPH_2:.*]]:2 = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<1024x3>, #tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_0

  func.return
}

func.func private @write_chain_func(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {
  %cst = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
  %0 = "tf.XlaAllReduce"(%arg0, %cst) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<1xf32>, tensor<1x2xi32>) -> tensor<1xf32>
  return %0 : tensor<1xf32>

  // CHECK-NEXT:     %[[GROUP:.*]] = "tf.Const"()
  // CHECK-SAME:       [0, 1, 2, 3]
  // CHECK-NEXT:     %[[REDUCED:.*]] = "tf.XlaAllReduce"(%[[ID]], %[[GROUP]]) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x4xi32>) -> tensor<f32>
  // CHECK-NEXT:     "tf.OpA"(%[[REDUCED]]) : (tensor<f32>) -> ()
  tf_device.replicate {n = 4 : i32} {
    "tf_device.cluster"() ({
      "tf.OpA"(%arg0) : (tensor<f32>) -> ()
      tf_device.return