tf.raw_ops.ApproxTopK

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

tf.raw_ops.ApproxTopK(
    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).