func.return %0#0 : tensor<*xi32>
}

func.func private @compute_0(%arg0: tensor<?x?xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
  func.return %arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>
}

func.func private @batch_3(%arg0: tensor<?x1xi32>) -> tensor<*xi32> {

  }

  // CHECK-LABEL: add_v2
  func.func @add_v2(%arg0: tensor<2xi32>) -> tensor<2xi32> {
    // CHECK: mhlo.add %arg0, %arg0 : tensor<2xi32>
    %0 = "tf.AddV2"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
    func.return %0: tensor<2xi32>
  }

  // CHECK-LABEL: not_allowlisted_op
  func.func @not_allowlisted_op(%arg0: tensor<3xi32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<?x?x?xf32> {

// CHECK-LABEL: @add_f32(%arg0: tensor<3x3xbf16>, %arg1: tensor<3x3xbf16>) -> tensor<3x3xbf16>
func.func @add_f32(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> {
  // CHECK-NOT: f32
  // CHECK: stablehlo.add
  %0 = stablehlo.add %arg0, %arg1: (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
  return %0 : tensor<3x3xf32>
}

// -----

func.func @avoidPackToConcatenationOnUnknownRank(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.pack"

// -----

func.func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {

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

// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1}>
// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu_fn_1}
// CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])

//go:build asan

#include "go_asm.h"
#include "textflag.h"

#define RARG0 R0
#define RARG1 R1
#define RARG2 R2
#define RARG3 R3
#define FARG R4

// Called from instrumented code.
// func runtime·doasanread(addr unsafe.Pointer, sz, sp, pc uintptr)
TEXT	runtime·doasanread(SB), NOSPLIT, $0-32
	MOVD	addr+0(FP), RARG0
	MOVD	sz+8(FP), RARG1
	MOVD	sp+16(FP), RARG2
	MOVD	pc+24(FP), RARG3

  func.return %0 : tensor<1x384x32xf32>
}

// PARTIAL-LABEL: @arg_reuse_test_2
func.func @arg_reuse_test_2(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func2"} {
  %0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>

  func.return %1 : tensor<1xf32>
}

func.func @cond_true(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

func.func @cond_false(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = tfl.mul %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>

    // non_tpu_ops
    %res_n = "tf.Add"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>

    return %res_n : tensor<i32>
  }
  func.func private @while_cond(%arg0: tensor<i32>) -> tensor<i1> {
    %0 = "tf.Less"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    return %0 : tensor<i1>
  }