module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<128x10xf32> {mhlo.sharding = "\08\03\1A\02\01\02\22\02\00\01"}, %arg1: tensor<10x1024xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg2: tensor<128x1024xf32> {mhlo.sharding = ""}) {
    func.return
  }
}

  // FB-EMPTY:
  // FB-NEXT:             }
  // FB-NEXT:       } ],

  // CHECK-LABEL: func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>)
  // CHECK-SAME:      -> tensor<40x40xf32>
  // CHECK:         %[[stat:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense
  // CHECK-SAME:      <[-1.000000e+00, 1.000000e+00]> : tensor<2xf32>}>
  // CHECK-SAME:      : (tensor<40x37xf32>) -> tensor<40x37xf32>

}

func.func @matmul_batch(%arg0: tensor<10x15xf32>, %arg1: tensor<15x17xf32>) -> tensor<10x17xf32> {
  %0 = "tf.BatchMatMul"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", device = "/device:CPU:0", name = "MatMul", adj_x = false, adj_y = false} :
(tensor<10x15xf32>, tensor<15x17xf32>) -> tensor<10x17xf32>
  func.return %0 : tensor<10x17xf32>
// CHECK-LABEL: matmul_batch

  // CHECK-DAG: %[[LOG_SOFTMAX_LOG:.*]] = "tf.Log"(%[[LOG_SOFTMAX_SUM]]) : (tensor<2x1xf32>) -> tensor<2x1xf32>
  // CHECK-DAG: %[[LOG_SOFTMAX:.*]] = "tf.Sub"(%[[LOG_SOFTMAX_SHIFTED]], %[[LOG_SOFTMAX_LOG]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32>

// CHECK: return %[[identity]] : tensor<1x3x2x2xf32>

func.func @cast_bf16_matmul_to_fp32(%arg0: tensor<1x10xf32>) -> (tensor<1x2xf32>) {
  %cst = "tf.Const"() {device = "", value = dense<1.000000e+01> : tensor<10x2xbf16>} : () -> tensor<10x2xbf16>
  %0 = "tf.Cast"(%arg0) {Truncate = false, device = ""} : (tensor<1x10xf32>) -> tensor<1x10xbf16>

  %11 = "mhlo.dot"(%10, %5) : (tensor<32x3920xf32>, tensor<3920x10xf32>) -> tensor<32x10xf32>
  %12 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<10xf32>) -> tensor<32x10xf32>
  %13 = mhlo.add %11, %12 : tensor<32x10xf32>
  %14 = mhlo.maximum %13, %0 : tensor<32x10xf32>
  return %14 : tensor<32x10xf32>
}
// CHECK: ReadVariableOp
// CHECK: mhlo.convolution
// CHECK: mhlo.add

  %2 = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
  // CHECK: tac.device = "CPU", tac.inference_type = "FLOAT"
  %3 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return
}

func.func @notAnnotateConst(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK-NOT: tac.device tac.inference_type

}

// -----

func.func @fully_connected_no_bias(%arg0: tensor<2x2x10xf32>, %arg1: tensor<40x40xf32>, %arg2: none) -> tensor<1x40xf32> {
  %0 = "tfl.fully_connected"(%arg0, %arg1, %arg2) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2x10xf32>, tensor<40x40xf32>, none) -> tensor<1x40xf32>
  func.return %0 : tensor<1x40xf32>
}

// -----

    // CHECK-NOT: tf.Multinomial
    %samples = "tf.Const"() { value = dense<10> : tensor<i32> } : () -> tensor<i32>
    %1 = "tf.Multinomial"(%arg0, %samples) {seed = 0, seed2 = 0}: (tensor<2x4xf32>, tensor<i32>) -> tensor<2x10xi32>
    func.return %1 : tensor<2x10xi32>
  }

  // CHECK-LABEL: @set_dynamic_dimension_size

  %max_size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NOT: tf.EmptyTensorList
  // CHECK: tf.BroadcastTo
  // CHECK-SAME: tensor<10x16xf32>
  %tl0 = "tf.EmptyTensorList"(%element_shape, %max_size) : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<16xf32>>>
  // CHECK-NOT: tf.TensorListReserve
  // CHECK: tf.BroadcastTo