Crate caffe2op_topk

source ·

Structs

  • | Given two tensors: X and K, retrieve | the top | | K[…, 1] elements from X on the last | dimension. | | X is an input tensor of shape [a_1, a_2, | …, a_n, r]. | | K is an input tensor of shape [a_1, a_2, | …, a_n, 1], where for each element, | r >= K[…, 1] > 0 | | Output two outputs: | | -Flatten values tensor of shape [ \sum_i | K[i, 1] ] which contains the values of | the top K[…, 1] elements along the | last dimension | | -Flatten indices tensor of shape [ \sum_i | K[i, 1] ] which contains the indices | of the top K[…, 1] elements, flatten | indices from the input tensor). | | These two outputs should be used with | the input K, so that we know which indices | in X are picked. | | Given two equivalent values, this operator | uses the indices along the last dim- | ension as a tiebreaker. That is, the | element with the lower index will appear | first. |
  • | Retrieve the top-K elements of the last | dimension. | | Given an input tensor of shape $(a_1, | a_2, …, a_n, r)$. k can be passed | as an integer argument or a 1D tensor | containing a single integer. | | Returns up to three outputs: | | 1. Value tensor of shape $(a_1, a_2, | …, a_n, k)$ which contains the values | of the top k elements along the last dimension | | 2. Index tensor of shape $(a_1, a_2, | …, a_n, k)$ which contains the indices | of the top k elements (original indices | from the input tensor). | | 3. [OPTIONAL] Flattened index tensor | of shape $(a_1 * a_2 * … * a_n * k,)$. | | Given two equivalent values, this operator | uses the indices along the last dimension | as a tiebreaker. | | That is, the element with the lower index | will appear first. | | Github Links: | | - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/top_k.cc |

Functions