// CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, 5, 21]> : tensor<3xi64>
  // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[21, 4]> : tensor<2xi64>
  // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg0, %[[cst]]) : (tensor<2x5x7x3xf32>, tensor<3xi64>) -> tensor<2x5x21xf32>
  // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg1, %[[cst_1]]) : (tensor<7x3x4xf32>, tensor<2xi64>) -> tensor<21x4xf32>

  }

  // 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:           [[VAL_62:%.*]] = "tf.Slice"([[VAL_53]], [[VAL_60]], [[VAL_61]]) : (tensor<4x3xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x0xf32>
// CHECK-DAG:       [[VAL_63:%.*]] = arith.constant dense<[2, 0]> : tensor<2xi64>
// CHECK-DAG:       [[VAL_64:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi64>
// CHECK:           [[VAL_65:%.*]] = "tf.Slice"([[VAL_53]], [[VAL_63]], [[VAL_64]]) : (tensor<4x3xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x0xf32>

  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

}

func.func @broadcast_args_i64(%arg0: tensor<3xi64>, %arg1: tensor<1xi64>) -> tensor<3xi64> {
  %0 = "tf.BroadcastArgs"(%arg0, %arg1) : (tensor<3xi64>, tensor<1xi64>) -> tensor<3xi64>
  func.return %0 : tensor<3xi64>

// CHECK-LABEL:broadcast_args_i64
// CHECK:  "tfl.broadcast_args"(%arg0, %arg1) : (tensor<3xi64>, tensor<1xi64>) -> tensor<3xi64>
}

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

func.func @testBroadcastGradientArgI64() -> (tensor<2xi64>, tensor<0xi64>) {
  %s0 = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64>
  %s1 = "tf.Const"() {value = dense<[2, 4]> : tensor<2xi64>} : () -> tensor<2xi64>
  %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<0xi64>, tensor<2xi64>) -> (tensor<2xi64>, tensor<0xi64>)

    %16 = "tf.BroadcastArgs"(%11, %14) : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %17 = "tf.Concat"(%cst_2, %16, %12) : (tensor<i32>, tensor<1xi64>, tensor<2xi64>) -> tensor<3xi64>
    %18 = "tf.Concat"(%cst_2, %16, %15) : (tensor<i32>, tensor<1xi64>, tensor<2xi64>) -> tensor<3xi64>
    %19 = "tf.BroadcastTo"(%4, %17) : (tensor<2x?x?xi8>, tensor<3xi64>) -> tensor<2x?x?xi8>