tf.approx_top_k

Returns min/max k values and their indices of the input operand in an approximate manner.

tf.approx_top_k(
    input,
    k,
    reduction_dimension=-1,
    recall_target=0.95,
    is_max_k=True,
    reduction_input_size_override=-1,
    aggregate_to_topk=True,
    name=None
)

See https://arxiv.org/abs/2206.14286 for the algorithm details. This op is only optimized on TPU currently.

input A Tensor. Must be one of the following types: half, bfloat16, float32. Array to search. Must be at least 1-D of the floating type
k An int that is >= 0. Specifies the number of min/max-k.
reduction_dimension An optional int. Defaults to -1. Integer dimension along which to search. Default: -1.
recall_target An optional float. Defaults to 0.95. Recall target for the approximation. Range in (0,1]
is_max_k An optional bool. Defaults to True. When true, computes max-k; otherwise computes min-k.
reduction_input_size_override An optional int. Defaults to -1. When set to a positive value, it overrides the size determined by input[reduction_dim] for evaluating the recall. This option is useful when the given input is only a subset of the overall computation in SPMD or distributed pipelines, where the true input size cannot be deferred by the input shape.
aggregate_to_topk An optional bool. Defaults to True. When true, aggregates approximate results to top-k. When false, returns the approximate results. The number of the approximate results is implementation defined and is greater equals to the specified k.
name A name for the operation (optional).

A tuple of Tensor objects (values, indices).
values A Tensor. Has the same type as input.
indices A Tensor of type int32.