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
operation
values

Public static functions

AggregateToTopk(bool x)
IsMaxK(bool x)
RecallTarget(float x)
ReductionDimension(int64 x)
ReductionInputSizeOverride(int64 x)

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
)