tensorflow::ops::ApproxTopK

#include <nn_ops.h>

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

Summary

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

Args:

• scope: A Scope object
• input: Array to search. Must be at least 1-D of the floating type
• k: Specifies the number of min/max-k.

Optional attributes (see Attrs):

• reduction_dimension: Integer dimension along which to search. Default: -1.
• recall_target: Recall target for the approximation. Range in (0,1]
• is_max_k: When true, computes max-k; otherwise computes min-k.
• reduction_input_size_override: 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: 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.

Returns:

• Output values: The min/max k values along the reduction_dimension of the input operand. The dimension are the same as the input operand except for the reduction_dimension: when aggregate_to_topk is true, the reduction dimension is k; otherwise, it is greater equals to k where the size is implementation-defined.
• Output indices: The indices of values along the reduction_dimension of the input operand.

Public attributes

::tensorflow::Output indices

Operation operation

::tensorflow::Output values

Public functions

 ApproxTopK(
  const ::tensorflow::Scope & scope,
  ::tensorflow::Input input,
  int64 k
)

 ApproxTopK(
  const ::tensorflow::Scope & scope,
  ::tensorflow::Input input,
  int64 k,
  const ApproxTopK::Attrs & attrs
)

Public static functions

Attrs AggregateToTopk(
  bool x
)

Attrs IsMaxK(
  bool x
)

Attrs RecallTarget(
  float x
)

Attrs ReductionDimension(
  int64 x
)

Attrs ReductionInputSizeOverride(
  int64 x
)