public final class
ApproxTopK
Returns min/max k values and their indices of the input operand in an approximate manner.
See https://arxiv.org/abs/2206.14286 for the algorithm details. This op is only optimized on TPU currently.
Nested Classes
| class | ApproxTopK.Options | Optional attributes for ApproxTopK
|
|
Public Methods
| static ApproxTopK.Options |
aggregateToTopk(Boolean aggregateToTopk)
|
| static <T extends Number> ApproxTopK<T> |
create(Scope scope, Operand<T> input, Long k, Options... options)
Factory method to create a class wrapping a new ApproxTopK operation.
|
| Output<Integer> |
indices()
The indices of `values` along the `reduction_dimension` of the `input` operand.
|
| static ApproxTopK.Options |
isMaxK(Boolean isMaxK)
|
| static ApproxTopK.Options |
recallTarget(Float recallTarget)
|
| static ApproxTopK.Options |
reductionDimension(Long reductionDimension)
|
| static ApproxTopK.Options |
reductionInputSizeOverride(Long reductionInputSizeOverride)
|
| Output<T> |
values()
The min/max k values along the `reduction_dimension` of the `input` operand.
|
Inherited Methods
| boolean |
equals(Object arg0)
|
| final Class<?> |
getClass()
|
| int |
hashCode()
|
| final void |
notify()
|
| final void |
notifyAll()
|
| String |
toString()
|
| final void |
wait(long arg0, int arg1)
|
| final void |
wait(long arg0)
|
| final void |
wait()
|
Public Methods
public static ApproxTopK.Options aggregateToTopk (Boolean aggregateToTopk)
Parameters
| aggregateToTopk | 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`. |
|---|
public static ApproxTopK<T> create (Scope scope, Operand<T> input, Long k, Options... options)
Factory method to create a class wrapping a new ApproxTopK operation.
Parameters
| scope | current scope |
|---|---|
| input | Array to search. Must be at least 1-D of the floating type |
| k | Specifies the number of min/max-k. |
| options | carries optional attributes values |
Returns
- a new instance of ApproxTopK
public Output<Integer> indices ()
The indices of `values` along the `reduction_dimension` of the `input` operand.
public static ApproxTopK.Options isMaxK (Boolean isMaxK)
Parameters
| isMaxK | When true, computes max-k; otherwise computes min-k. |
|---|
public static ApproxTopK.Options recallTarget (Float recallTarget)
Parameters
| recallTarget | Recall target for the approximation. Range in (0,1] |
|---|
public static ApproxTopK.Options reductionDimension (Long reductionDimension)
Parameters
| reductionDimension | Integer dimension along which to search. Default: -1. |
|---|
public static ApproxTopK.Options reductionInputSizeOverride (Long reductionInputSizeOverride)
Parameters
| reductionInputSizeOverride | 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. |
|---|
public Output<T> 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.