%1 = "tf.OpB"() {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\02"} : () -> tensor<2x2xi64>
        // expected-error @+1 {{should have the same sharding}}
        %2 = "tf.OpC"(%0, %1) {_xla_map_outside_compilation = "0", _xla_outside_compilation = "from_launch"} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
        "tf.OpD"(%2) : (tensor<2x2xi64>) -> ()
        tf_device.return

  %0 = "tf.SpaceToBatchND"(%input, %block_shape, %paddings) : (tensor<3x5x7x10xf32>, tensor<2xi64>, tensor<2x2xi64>) -> tensor<36x2x3x10xf32>
  func.return %0 : tensor<36x2x3x10xf32>
}

// -----

// Test invalid tf.SpaceToBatchND

}

func.func @matrix_set_diag_v3_non_zero_k(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi32>) -> tensor<3x3xi64> {
  %cst = arith.constant dense<1> : tensor<i32>
  %0 = "tf.MatrixSetDiagV3"(%arg0, %arg1, %cst) : (tensor<3x3xi64>, tensor<3xi32>, tensor<i32>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

// CHECK-LABEL: @matrix_set_diag_v3_non_zero_k
// CHECK: tf.MatrixSetDiagV3
}

  // CHECK-SAME: (tensor<3x3xi64>, tensor<3xi64>, tensor<i32>) -> tensor<3x3xi64>
}

// CHECK-LABEL: @testMatrixSetDiagV2
func.func @testMatrixSetDiagV2(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi64>, %arg2: tensor<i32>) -> tensor<3x3xi64> {
  %0 = "tf.MatrixSetDiagV2"(%arg0, %arg1, %arg2) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x3xi64>, tensor<3xi64>, tensor<i32>) -> tensor<3x3xi64>

// CHECK-LABEL: testAddWithI64Broadcasting
func.func @testAddWithI64Broadcasting(tensor< 2x3xi64>, tensor<3xi64>) -> tensor<2x3xi64> {
^bb0(%arg0: tensor<2x3xi64>, %arg1: tensor<3xi64>):
  // CHECK: tfl.add(%arg0, %arg1)
  %0 = tfl.add(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor< 2x3xi64>, tensor<3xi64>) -> tensor<2x3xi64>
  func.return %0#0 : tensor<2x3xi64>
}

// -----

// CHECK-LABEL: add_with_i32_five_dim_broadcasting

    %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi64>} : () -> tensor<3xi64>
    // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<8x4x14x14x16xf32>

// For example,
//
//   %0 = tf.Pack(%input) {axis = 0} // %input : tensor<2x3xf32>
//
// can be canonicalized to
//
//   %shape = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi64>}
//   %0 = tf.Reshape(%input, %shape)
struct ConvertPackToReshape : public OpRewritePattern<PackOp> {
  using OpRewritePattern<PackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(PackOp pack_op,