%1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_1} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_1} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

  // MLIR @main function corresponds to the TF function "main_original".
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
      return %arg : tensor<1x2xf32>
    }
  )mlir");
  ASSERT_TRUE(module_op);

  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases;

    return %1 : tensor<1x3xf32>
  }

  func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    return %0 : tensor<1x3xf32>

}

// CHECK-LABEL: @broadcast_batch_matmul_v2_rhs
func.func @broadcast_batch_matmul_v2_rhs(%arg0: tensor<17x17x17xf32>, %arg1: tensor<17x24xf32>) -> tensor<17x17x24xf32> {
  %cst = arith.constant dense<[17, 17, 24]> : tensor<3xi64>
  %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<17x24xf32>, tensor<3xi64>) -> tensor<17x17x24xf32>

  %1 = "tf.BiasAdd"(%0, %arg0) {data_format = "NHWC"} : (tensor<*xf32>, tensor<128xf32>) -> tensor<*xf32>
  %2 = "tf.Identity"(%1) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %2 : tensor<*xf32>
}

// CHECK-LABEL: conv2DBiasAdd_reluActivation
func.func @conv2DBiasAdd_reluActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) {

func.func @single_state_single_dataset_type_multiple_arguments(
      %arg0: tensor<!tf_type.variant>,
      %arg1: tensor<i64>,
      %arg2: tensor<!tf_type.resource<tensor<64xf32>>>,
      %arg3: tensor<!tf_type.resource<tensor<128xf32>>>
    ) {
    // CHECK:      %[[ANON_ITER:[0-9]*]] = "tf.AnonymousIteratorV3"

func.func @tfl_wrapped_jax_random_uniform(%arg0: tensor<2xui32>) -> tuple<tensor<1x2xf32>> {
  // This is a fake jax random uniform body.
  %0 = stablehlo.constant dense<0.0> : tensor<2xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<2xf32>) -> tensor<1x2xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<1x2xf32>) -> tuple<tensor<1x2xf32>>
  func.return %2 : tuple<tensor<1x2xf32>>

  %dq_w = "tfl.dequantize"(%w) : (tensor<12x2x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>) -> tensor<12x2xf32>
  %emb = "tfl.gather"(%dq_w, %arg0) {axis = 0 : i32, batch_dims = 0 : i32} : (tensor<12x2xf32>, tensor<2xi32>) -> tensor<2x2xf32>
  func.return %emb : tensor<2x2xf32>

// CHECK-NEXT: tfl.pseudo_qconst
// CHECK-NEXT: tfl.gather
// CHECK-NEXT: tfl.dequantize
}

  // CHECK: "tfl.conv_2d"(%arg0, %arg1, %arg2) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x112x80x128xf32>, tensor<128x3x3x128xf32>, tensor<128xf32>) -> tensor<1x108x76x128xf32>