// CHECK-LABEL: doubleTposeInput
func.func @doubleTposeInput(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
  %perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
  %0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
  %perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
  %1 = "tfl.transpose"(%arg1, %perm1) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>

  func.func private @if_then_branch(%arg0: tensor<*xi1>, %arg1: tensor<*xi64>, %arg2: tensor<*xi64>) -> tensor<*xi1> {
    func.return %arg0 : tensor<*xi1>
  }
  func.func private @if_else_branch(%arg0: tensor<*xi1>, %arg1: tensor<*xi64>, %arg2: tensor<*xi64>) -> tensor<*xi1> {
    func.return %arg0 : tensor<*xi1>
  }
  func.func @"tpu_subgraph"(%arg0: tensor<*xi64>) -> tensor<*xi1> {

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s

func.func @main(%arg0: tensor<32x4x4x128xf32>, %arg1: tensor<1x32x42x128xf32>, %arg2: tensor<4xi32>) -> tensor<1x64x84x32xf32> {

func.func @main(%arg0: tensor<32x4x4x128xf32>, %arg1: tensor<1x32x42x128xf32>, %arg2: tensor<4xi32>) -> tensor<1x64x84x32xf32> {

module {
  // CHECK-LABEL: main
  func.func @main(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
    %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
    func.return %0 : tensor<4xf32>
    // CHECK:  [[VAL_0:%.*]] = tfl.sub %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<4xf32>

func (m *MockRuntimeCache) ForceUpdateIfOlder(arg0 context.Context, arg1 time.Time) error {
	m.ctrl.T.Helper()
	ret := m.ctrl.Call(m, "ForceUpdateIfOlder", arg0, arg1)
	ret0, _ := ret[0].(error)
	return ret0
}

// ForceUpdateIfOlder indicates an expected call of ForceUpdateIfOlder.
func (mr *MockRuntimeCacheMockRecorder) ForceUpdateIfOlder(arg0, arg1 any) *gomock.Call {
	mr.mock.ctrl.T.Helper()

// CHECK-LABEL: func @testPwStreamResults
func.func @testPwStreamResults(%arg0: tensor<f32>, %arg1: tensor<f32>) {
  "tf.PwStreamResults"(%arg0, %arg1) {names = ["foo", "bar"]} : (tensor<f32>, tensor<f32>) -> ()
  return
}

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {

// CHECK:           return [[RESULT]] : tensor<3x4x4xf32>
// CHECK:         }

  %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
  func.return %0 : tensor<3x4x4xf32>
}

func.func @batchmatmulv2_lhs_batch(%arg0: tensor<3x4x2xf32>, %arg1: tensor<2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL:   func @batchmatmulv2_lhs_batch

// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<3x2xf32>
func.func @convert_quantfork_qdq_to_stablehlo_uniform_qdq(%arg0: tensor<1x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<1x2xf32> {
  // CHECK: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]]
  // CHECK-NOT: "quantfork.qcast"
  // CHECK: %[[Q2:.*]] = stablehlo.uniform_quantize %[[ARG1]]
  // CHECK-NOT: "quantfork.qcast"

  return
}

// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>) -> () {
  return
}
}

// -----
// Multiple ifrt calls and have two sharded arguments

// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []