Caffe2 - C++ API
A deep learning, cross platform ML framework
partition_ops.cc
1 
17 #include "caffe2/operators/partition_ops.h"
18 
19 namespace caffe2 {
20 namespace {
21 
22 REGISTER_CPU_OPERATOR(Partition, PartitionOp);
23 REGISTER_CPU_OPERATOR(LengthsPartition, LengthsPartitionOp);
24 REGISTER_CPU_OPERATOR(GatherByKey, GatherByKeyOp);
25 
26 OPERATOR_SCHEMA(GatherByKey)
27  .NumInputs(2, INT_MAX)
28  .NumOutputs(1)
29  .SetDoc(R"DOC(
30 Inverse operation of Partition.
31 
32 Takes the original, full 'keys' tensor followed by sharded value tensors,
33 and returns the full value tensor, combined using the same hash used in
34 Partition.
35 )DOC")
36  .Input(
37  0,
38  "keys",
39  "The first input is the full keys tensor"
40  " (same as the first input of Partition).")
41  .Input(
42  1,
43  "sharded_values",
44  "Subsequented inputs are sharded values tensors.")
45  .Output(0, "values", "Reconstructed values tensor.");
46 
47 OPERATOR_SCHEMA(Partition)
48  .NumInputsOutputs([](int in, int out) {
49  return in > 0 && out > 0 && out % in == 0;
50  })
51  .SetDoc(R"DOC(
52 Splits the input int tensor into multiple ones according to the first tensor.
53 
54 Takes the first input and partitions it to shards according to the remainder of
55 values modulo the number of partitions. It requires that the first tensor is of
56 integral type. The number of partitions is derived as (num_output / num_input).
57 
58 If additional inputs are present they must have the same shape as the first
59 input, optionally with extra trailing dimensions. They will be partitioned
60 accordingly to the first input.
61 
62 Optional arg 'pack_first_input' transforms the first tensor values as
63 X_ij / num_partitions.
64 
65 Outputs are ordered as
66 X_0_part_0, X_1_part_0, ..., X_N-1_part_0, X_0_part_1, ..., X_N-1_part_K-1
67 )DOC")
68  .Arg(
69  "pack_first_input",
70  "(int, default 0) If set, the operator transforms "
71  "the first tensor values as floor(X_ij / num_partitions)")
72  .Input(
73  0,
74  "input",
75  "Input tensor containing data to be partitioned. The "
76  "number of input tensors might be greater than 1 but must have the "
77  "same shape as the previous tensors.")
78  .Output(
79  0,
80  "partitions",
81  "Output Partitions. The number of output tensors has to be a "
82  "multiple of the number of input tensors.");
83 
84 OPERATOR_SCHEMA(LengthsPartition)
85  .NumInputsOutputs([](int in, int out) {
86  return in >= 2 && out > 0 && out % in == 0;
87  })
88  .SetDoc(R"DOC(
89 LengthsPartition splits the input int tensor into multiple ones according to the
90 second tensor. The first dimension is expected to be the tensor that describes
91 lengths of the elements.
92 
93 Takes the second input and partitions it to shards according to the remainder of
94 values modulo the number of partitions. It requires the second tensor to be
95 a 1D-tensor of the integral type. The first tensor should be 1D-tensor of int32
96 that would represent the lengths of the elements in the input. The number of
97 partitions is derived as (num_output / num_input).
98 
99 If additional inputs are present they must have the same shape as the first
100 input, optionally with extra trailing dimensions. They will be partitioned
101 accordingly to the first input.
102 
103 Optional arg 'pack_first_input' transforms the first tensor values as
104 X_ij / num_partitions.
105 
106 Outputs are ordered as
107 X_0_part_0, X_1_part_0, ..., X_N-1_part_0, X_0_part_1, ..., X_N-1_part_K-1
108 )DOC")
109  .Arg(
110  "pack_first_input",
111  "(int, default 0) If set, the operator transforms "
112  "the first tensor values as floor(X_ij / num_partitions)")
113  .Input(
114  0,
115  "input",
116  "Input tensor containing data to be partitioned. The "
117  "number of input tensors might be greater than 1 but must have the "
118  "same shape as the previous tensors.")
119  .Output(
120  0,
121  "partitions",
122  "Output Partitions. The number of output tensors has to be a "
123  "multiple of the number of input tensors.");
124 
125 // This should actually have gradient, but for now nothing uses it.
126 // Because gradient computation right now is not input/output aware it can't be
127 // GRADIENT_NOT_IMPLEMENTEDYET
128 NO_GRADIENT(Partition);
129 NO_GRADIENT(LengthsPartition);
130 } // namespace
131 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.