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.

Constructors and Destructors
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 attributes
indices	::tensorflow::Output
operation	Operation
values	::tensorflow::Output

Public static functions
AggregateToTopk(bool x)	Attrs
IsMaxK(bool x)	Attrs
RecallTarget(float x)	Attrs
ReductionDimension(int64 x)	Attrs
ReductionInputSizeOverride(int64 x)	Attrs

Structs
tensorflow::ops::ApproxTopK::Attrs	Optional attribute setters for ApproxTopK.

Public attributes

indices

::tensorflow::Output indices

operation

Operation operation

values

::tensorflow::Output values

Public functions

ApproxTopK

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

ApproxTopK

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

Public static functions

AggregateToTopk

Attrs AggregateToTopk(
  bool x
)

IsMaxK

Attrs IsMaxK(
  bool x
)

RecallTarget

Attrs RecallTarget(
  float x
)

ReductionDimension

Attrs ReductionDimension(
  int64 x
)

ReductionInputSizeOverride

Attrs ReductionInputSizeOverride(
  int64 x
)