} : (tensor<2x7xi64>) -> (tensor<2x7xi32>)
  %1 = "tf.Identity"(%0) : (tensor<2x7xi32>) -> tensor<2x7xi32>
  func.return %1 : tensor<2x7xi32>
}

// CHECK-LABEL: func @func
// CHECK-SAME: %arg0: tensor<2x7xi64> {mhlo.sharding = "\08\03\1A\03\01\01\02\22\02\00\01B\01\00"
func.func @func(%arg0: tensor<2x7xi64>) -> (tensor<2x7xi32>) {

  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
    %7 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
    stablehlo.return %7 : tensor<f32>
  }) {padding = dense<[[0, 0], [0, 1], [0, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>, window_strides = array<i64: 1, 2, 2, 1>} : (tensor<?x112x112x64xf32>, tensor<f32>) -> tensor<?x56x56x64xf32>
  return %6 : tensor<?x56x56x64xf32>

// CHECK: %[[id:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id]]) <{shape = #tf_type.shape<1x2x3x4>}> : (tensor<*x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
// CHECK: return %[[es]] : tensor<1x2x3x4x!tf_type.string>
}

// CHECK-LABEL: decompose_tf_consecutive

  func.func @main() -> () {
    %cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
    %0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
    func.return
  }
})";

absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
    const char* mlir_module_str,
    ConfigProto::Experimental::MlirBridgeRollout rollout_state) {

      stablehlo.return %3 : tensor<f32>
  }) {
    window_dimensions = array<i64: 1, 1, 2, 2>,
    window_strides = array<i64: 1, 1, 2, 2>,
    padding = dense<[[0, 0], [0, 0], [1, 1], [1, 1]]> : tensor<4x2xi64>
  } : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x9x9xf32>
  return %2 : tensor<1x4x9x9xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: stablehlo.constant

  // CHECK: "tf.Conv3DBackpropInputV2"
}

func.func @mul_i64(%arg0: tensor<14xi64>, %arg1: tensor<14xi64>) -> tensor<14xi64> {
  %0 = "tf.Mul"(%arg0, %arg1) : (tensor<14xi64>, tensor<14xi64>) -> tensor<14xi64>
  func.return %0: tensor<14xi64>

// CHECK-LABEL: mul_i64
// CHECK:  tfl.mul %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<14xi64>
// CHECK:  return
}

func.func @reshape_vector_shape(tensor<4x4x4xf32>) -> tensor<16x4xf32> {
^bb0(%arg0: tensor<4x4x4xf32>) :
  %shape0 = arith.constant dense<[[16, 4]]> : tensor<1x2xi32>
  // expected-error @+1 {{'tfl.reshape' op requires 'shape' to be rank 1, but got 2}}
  %1 = "tfl.reshape"(%arg0, %shape0) : (tensor<4x4x4xf32>, tensor<1x2xi32>) -> tensor<16x4xf32>
  func.return %1 : tensor<16x4xf32>
}

// -----

  // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSH_VAL]], %[[UPDATE_SHAPE]]) : (tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32>
  // CHECK-NEXT: %[[ZERO_INDS:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
  // CHECK-NEXT: %[[CONCAT_DIM:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

  }

  // CHECK-LABEL: func @xla_reduce_window
  func.func @xla_reduce_window(%arg0: tensor<7xf32>, %arg1: tensor<f32>) -> tensor<?xf32> {
    %cst = "tf.Const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
    %cst_0 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
    %cst_1 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>

    %cell_input = "tfl.pseudo_qconst"() {qtype = tensor<1x20x!quant.uniform<i16:f32, 0.006:-34>>, value = dense<1> : tensor<1x20xi6>} : () -> tensor<1x20x!quant.uniform<i16:f32, 0.006:-34>>
    %0 = "tfl.unidirectional_sequence_lstm"(%input,
      %cst_11, %cst_11, %cst_11, %cst_11,
      %cst_3, %cst_3, %cst_3, %cst_3,
      %cst_2, %cst_2, %cst_2,