// CHECK: %[[arg2_i64:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2xi64>
// CHECK: %[[slice:.*]] = "tf.Slice"(%arg0, %[[tensor_scatter_update]], %[[arg2_i64]]) : (tensor<?x2xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<*xf32>
// CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[slice]], %[[cst_1]]) : (tensor<*xf32>, tensor<1xi64>) -> tensor<*xf32>

  %cst = arith.constant dense<[2, 100]> : tensor<2xi64>
  // CHECK: tac.cost = 4.000000e+01
  %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
  // CHECK: tac.cost = 4.040000e+01
  %1 = "tfl.reshape"(%0, %cst) {tac.device = "GPU"} : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
  func.return %1 : tensor<2x100xf32>
}

  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>
    func.return %0 : tensor<?xi64>
  }

// RUN: tf-opt %s -tfl-prepare-tf -tfl-legalize-tf='run-tfl-runtime-verification=false'  -tfl-optimize | FileCheck %s

func.func @broadcast_to_bf16(%arg0: tensor<3xbf16>, %arg1: tensor<2xi64>) -> tensor<3x3xbf16> {
  %0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xbf16>, tensor<2xi64>) -> tensor<3x3xbf16>
  func.return %0: tensor<3x3xbf16>

// CHECK-LABEL: broadcast_to_bf16
// CHECK:  [[CST:%.*]] = arith.constant dense<1.000000e+00> : tensor<3x3xbf16>

func.func @broadcast_both_operand(%arg0: tensor<7xf32>, %arg1: tensor<5x1xf32>) -> tensor<5x7xf32> {
  %cst = arith.constant dense<[5, 7]> : tensor<2xi64>
  %0 = "tf.BroadcastTo"(%arg0, %cst) : (tensor<7xf32>, tensor<2xi64>) -> tensor<5x7xf32>
  %1 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<5x1xf32>, tensor<2xi64>) -> tensor<5x7xf32>
  %2 = "tf.Add"(%0, %1) : (tensor<5x7xf32>, tensor<5x7xf32>) -> tensor<5x7xf32>
  func.return %2 : tensor<5x7xf32>

  // CHECK-DAG: %cst = arith.constant dense<0.000000e+00> : tensor<10x10xf32>
  // CHECK-DAG: %cst_0 = arith.constant dense<1000> : tensor<2xi64>
  // CHECK-DAG: %cst_1 = arith.constant dense<1.000000e+00> : tensor<f32>
  // CHECK: %0 = "tfl.fill"(%cst_0, %cst_1) : (tensor<2xi64>, tensor<f32>) -> tensor<1000x1000xf32>
  // CHECK: return %cst, %0 : tensor<10x10xf32>, tensor<1000x1000xf32>

  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<10x19xf32> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>
    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<10x19xf32>
    func.return %1 : tensor<10x19xf32>
  }
}

// Tests that foldable ops are constant-folded to enable legalization of ops

// 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]

  // 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]])

  %cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
  // CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%[[CONST]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  // CHECK: %[[MUL:.*]] = mhlo.multiply %[[BROADCAST]], %[[ARG]] : tensor<1x1x2x4xi32>