// CHECK-LABEL:   func @batchmatmulv2_lhs_batch
// CHECK:           "mhlo.dynamic_broadcast_in_dim"({{.*}}, {{.*}}) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}>
// CHECK:           "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME:        lhs_batching_dimensions = [0]
// CHECK-SAME:        rhs_batching_dimensions = [0]
// CHECK-SAME:        lhs_contracting_dimensions = [2]

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

      // CHECK:      name: "ParseExample"
      // CHECK-NEXT: op: "ParseExampleV2"
      // CHECK-NEXT: input: "input0"
      // CHECK-NEXT: input: "tf.Const{{_.*_3}}"

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 3, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

  // CHECK-LABEL: func @init_all_tables_with_explicit_vocab_size
  // CHECK-DAG: [[CST:%.*]]  = arith.constant dense<["apple", "banana"]> : tensor<2x!tf_type.string>
  // CHECK-DAG: [[CST_0:%.*]]  = arith.constant dense<[0, 1]> : tensor<2xi64>
  // CHECK: [[VAL:%.*]] = "tf.HashTableV2"()
  // CHECK: "tf.LookupTableImportV2"([[VAL]], [[CST]], [[CST_0]])

  // CHECK-LABEL: func @init_all_tables_with_explicit_vocab_size
  // CHECK-DAG: [[CST:%.*]]  = arith.constant dense<["apple", "banana"]> : tensor<2x!tf_type.string>
  // CHECK-DAG: [[CST_0:%.*]]  = arith.constant dense<[0, 1]> : tensor<2xi64>
  // CHECK: [[VAL:%.*]] = "tf.HashTableV2"()
  // CHECK: "tf.LookupTableImportV2"([[VAL]], [[CST]], [[CST_0]])

    %24 = "tf.Softmax"(%18) : (tensor<2x1001xf32>) -> tensor<2x1001xf32>
    %25 = "tf.ArgMax"(%24, %2) : (tensor<2x1001xf32>, tensor<i32>) -> tensor<2xi64>
    %26 = "tf.Cast"(%25) {Truncate = false} : (tensor<2xi64>) -> tensor<2xf32>
    %27 = "tf.Equal"(%14, %26) {incompatible_shape_error = true} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>

  %cst1 = mhlo.constant dense<[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]]> : tensor<1x2x2x3xf32>
  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<[1, 3]> : tensor<2xi64>}> : (tensor<2x3xf32>) -> tensor<1x2x2x3xf32>
  %1 = mhlo.multiply %0, %cst1 : tensor<1x2x2x3xf32>
  // CHECK:      return %[[RES]] : tensor<1x2x2x3xf32>
  func.return %1 : tensor<1x2x2x3xf32>
}

// Test models with dynamic bounds ops.

// CHECK-LABEL: func private @test_fuse_dynamic_dimension_ops
func.func private @test_fuse_dynamic_dimension_ops(%arg0: tensor<?x?xi32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<2xi64>, %arg3: tensor<?xi64>, %arg4: tensor<?xi64>) -> tensor<*xi32> {
  // CHECK-NOT: tf._TPUCompileMlirOp
  // CHECK-NOT: tf.TPUCompileSucceededAssert
  // CHECK-NOT: tf.TPUExecuteOp
  // CHECK-NOT: tf.SetStaticDimensionBounds

// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32>
// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() <{value = dense<[1024, 3]> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64>
// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>