func.func @main(%arg0: tensor<3x3x10xbf16>, %arg1: tensor<3xi32>) -> tensor<1x?x4xbf16> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2", outputs = "_retval0"}} {
      %cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>

    %2:2 = "tf.WhileRegion"(%1, %arg0) ({
      ^bb0(%carg0: tensor<f32>, %carg1: tensor<i32>):
         %limit = arith.constant dense<5> : tensor<i32>
         %cond = "tf.NotEqual"(%carg1, %limit) : (tensor<i32>, tensor<i32>) -> tensor<i1>
         "tf.Yield"(%cond) : (tensor<i1>) -> ()
    },  {
      ^bb0(%barg0: tensor<f32>, %barg1: tensor<i32>):

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
  // CHECK: return %arg0
  func.return %0 : tensor<?x?xf32>
}

// -----

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<*xf32>
  // CHECK: "tfl.squeeze"

		{name: "MODSD", argLength: 2, reg: gp21, asm: "MODSD", typ: "Int64"},  // arg0 % arg1 (signed 64-bit)
		{name: "MODUW", argLength: 2, reg: gp21, asm: "MODUW", typ: "UInt32"}, // arg0 % arg1 (unsigned 32-bit)
		{name: "MODSW", argLength: 2, reg: gp21, asm: "MODSW", typ: "Int32"},  // arg0 % arg1 (signed 32-bit)
		// MOD is implemented as rem := arg0 - (arg0/arg1) * arg1

  // CHECK: cond

  // CHECK: do
  // CHECK: mhlo.return %[[ARG0]], %[[ARG1]] : tensor<2x?xi8>, tensor<?x2xi8>

  %2:2 = "tf.WhileRegion"(%0, %1) ({
  ^bb0(%carg0: tensor<2x2x!tf_type.qint8>, %carg1: tensor<2x2x!tf_type.qint8>):
    %check = "tf.Equal"(%arg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "tf.Yield"(%check) : (tensor<i1>) -> ()
    }, {
  ^bb0(%barg0: tensor<2x?x!tf_type.qint8>, %barg1: tensor<?x2x!tf_type.qint8>):

	// can be executed on g0. Second, it is called frequently, so will
	// benefit from this fast path.
	CMPQ	RARG0, $0
	JNE	rest
	get_tls(RARG0)
	MOVQ	g(RARG0), RARG0
	MOVQ	g_m(RARG0), RARG0
	MOVQ	m_p(RARG0), RARG0
	MOVQ	p_raceprocctx(RARG0), RARG0
	MOVQ	RARG0, (RARG1)
	RET

rest:
	// Transition from C ABI to Go ABI.
	PUSH_REGS_HOST_TO_ABI0()
	// Set g = g0.
	get_tls(R12)

}

TEST(IsSupportedByReplicatedBridge, ReplicatedModule) {
  const char* const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
  %0 = "tf.Identity"(%arg0) {_tpu_replicate = "cluster"} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;

	{name: "Signmask", argLength: 1, typ: "Int32"},  // 0 if arg0 >= 0, -1 if arg0 < 0
	{name: "Zeromask", argLength: 1, typ: "UInt32"}, // 0 if arg0 == 0, 0xffffffff if arg0 != 0
	{name: "Slicemask", argLength: 1},               // 0 if arg0 == 0, -1 if arg0 > 0, undef if arg0<0. Type is native int size.

	{name: "SpectreIndex", argLength: 2},      // arg0 if 0 <= arg0 < arg1, 0 otherwise. Type is native int size.

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> {

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