*   Added reference implementation for 16-bit int and 32-bit unsigned int unquantized `mul`.
    *   `add_op` supports broadcasting up to 6 dimensions.
    *   Added 16-bit support for `top_k`.

*   `tf.function`

                               /*outfeed_config=*/rewriter.getStringAttr(""));
    rewriter.eraseOp(op);
    return success();
  }
};

// Converts tf.TopKV2 to chlo.top_k.
class ConvertTopKV2Op : public OpRewritePattern<TF::TopKV2Op> {
 public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::TopKV2Op op,

    backend_config = {
      aggregate_to_topk = true,
      is_fallback = true,
      recall_target = 8.500000e-01 : f32,
      reduction_dim = 1 : i64,
      reduction_input_size_override = -1 : i64,
      top_k = 4 : i64}
    } : (tensor<1x4xf32>, tensor<1x4xi32>, tensor<f32>, tensor<i32>) -> (tensor<1x4xf32>, tensor<1x4xi32>)
  func.return %0#0, %0#1 : tensor<1x4xf32>, tensor<1x4xi32>
}

func.func @topk_v2(%input: tensor<16x16xf32>) -> (tensor<16x8xf32>, tensor<16x8xi32>) {
  %k = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>

  // CHECK:     chlo.top_k(%[[INPUT]], k = 8)
  %0:2 = "tf.TopKV2"(%input, %k): (tensor<16x16xf32>, tensor<i32>) -> (tensor<16x8xf32>, tensor<16x8xi32>)
  func.return %0#0, %0#1: tensor<16x8xf32>, tensor<16x8xi32>
}

    checkNotNull(comparator);
    return Collector.of(
        () -> TopKSelector.<T>least(k, comparator),
        TopKSelector::offer,
        TopKSelector::combine,
        TopKSelector::topK,
        Collector.Characteristics.UNORDERED);
  }

  /**
   * Returns a {@code Collector} that returns the {@code k} greatest (relative to the specified

        {
          "datasource": {
            "type": "prometheus",
            "uid": "${DS_PROMETHEUS}"
          },
          "exemplar": true,
          "expr": "topk(1, sum(minio_cluster_capacity_usable_total_bytes{job=~\"$scrape_jobs\"}) by (instance)) - topk(1, sum(minio_cluster_capacity_usable_free_bytes{job=~\"$scrape_jobs\"}) by (instance))",
          "format": "time_series",
          "instant": false,
          "interval": "1m",

    checkNotNull(comparator);
    return Collector.of(
        () -> TopKSelector.<T>least(k, comparator),
        TopKSelector::offer,
        TopKSelector::combine,
        TopKSelector::topK,
        Collector.Characteristics.UNORDERED);
  }

  /**
   * Returns a {@code Collector} that returns the {@code k} greatest (relative to the specified

  // to be consistent with other entrypoints.
  pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());

  // Decompose CHLO into StableHLO ops
  // TODO(b/331843141): There are some CHLO's like TopK which we could instead
  // lower to TFL ops.
  mlir::stablehlo::experimental::createChloLegalizeToStablehloPipeline(
      pass_manager);
  pass_manager.addPass(mlir::odml::CreateTransposeCommuteOpsPass());

                       (TFL_MeanOp $arg0, $arg1, $arg2)>;

def LegalizeSum : Pat<(TF_SumOp $arg, $axes, BoolAttr:$arg2),
                      (TFL_SumOp $arg, (CreateTFCastToInt32Op $axes), $arg2)>;

// TopK in TFL is always sorted so we ignore that attribute here.
def LegalizeTopKV2 : Pat<(TF_TopKV2Op $input, $k, $ignored_sorted),
                         (TFL_TopKV2Op $input, $k)>;

  EXPECT_NEAR(expected[1], fn.result0(0, 1), 1e4);
  EXPECT_NEAR(expected[2], fn.result0(1, 0), 1e4);
  EXPECT_NEAR(expected[3], fn.result0(1, 1), 1e4);
}

TEST(TFCompileTest, TopK) {
  Eigen::ThreadPool tp(1);
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());

  TopKComp fn;

  fn.set_thread_pool(&device);
  // x = [4, 1, 4, 4, 3]
  fn.arg0(0) = 4;
  fn.arg0(1) = 1;