File: merge_id_lists_op.cc

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (36 lines) | stat: -rw-r--r-- 1,637 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#include "caffe2/operators/merge_id_lists_op.h"

namespace caffe2 {
namespace {
REGISTER_CPU_OPERATOR(MergeIdLists, MergeIdListsOp<CPUContext>);

OPERATOR_SCHEMA(MergeIdLists)
    .NumInputs([](int n) { return (n > 0 && n % 2 == 0); })
    .NumOutputs(2)
    .SetDoc(R"DOC(
MergeIdLists: Merge multiple ID_LISTs into a single ID_LIST.

An ID_LIST is a list of IDs (may be ints, often longs) that represents a single
feature. As described in https://caffe2.ai/docs/sparse-operations.html, a batch
of ID_LIST examples is represented as a pair of lengths and values where the
`lengths` (int32) segment the `values` or ids (int32/int64) into examples.

Given multiple inputs of the form lengths_0, values_0, lengths_1, values_1, ...
which correspond to lengths and values of ID_LISTs of different features, this
operator produces a merged ID_LIST that combines the ID_LIST features. The
final merged output is described by a lengths and values vector.

WARNING: The merge makes no guarantee about the relative order of ID_LISTs
within a batch. This can be an issue if ID_LIST are order sensitive.
)DOC")
    .Input(0, "lengths_0", "Lengths of the ID_LISTs batch for first feature")
    .Input(1, "values_0", "Values of the ID_LISTs batch for first feature")
    .Output(0, "merged_lengths", "Lengths of the merged ID_LISTs batch")
    .Output(1, "merged_values", "Values of the merged ID_LISTs batch");
NO_GRADIENT(MergeIdLists);
}
}
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
    MergeIdLists,
    "_caffe2::MergeIdLists(Tensor[] lengths_and_values) -> (Tensor merged_lengths, Tensor merged_values)",
    caffe2::MergeIdListsOp<caffe2::CPUContext>);