forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sparse_to_dense_mask_op.cc
138 lines (123 loc) · 4.79 KB
/
sparse_to_dense_mask_op.cc
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#include "caffe2/operators/sparse_to_dense_mask_op.h"
namespace caffe2 {
namespace {
REGISTER_CPU_OPERATOR(SparseToDenseMask, SparseToDenseMaskOp<CPUContext>);
REGISTER_CPU_OPERATOR(
SparseToDenseMaskGradient,
SparseToDenseMaskGradientOp<CPUContext>);
OPERATOR_SCHEMA(SparseToDenseMask)
.NumInputs(3, 4)
.NumOutputs(1, 2)
.DisallowInputFillers() // TODO: enable the filler
.TensorInferenceFunction([](const OperatorDef& def,
const vector<TensorShape>& in) {
ArgumentHelper helper(def);
auto mask = helper.template GetRepeatedArgument<int64_t>("mask");
bool return_presence_mask = helper.template GetSingleArgument<bool>(
"return_presence_mask", false);
vector<TensorShape> out(1);
if (in.size() == 4) {
out[0].add_dims(in[3].dims(0));
}
out[0].add_dims(mask.size());
for (const auto dim : in[2].dims()) {
out[0].add_dims(dim);
}
out[0].set_data_type(in[2].data_type());
if (return_presence_mask) {
out.emplace_back();
if (in.size() == 4) {
out[1].add_dims(in[3].dims(0));
}
out[1].add_dims(mask.size());
out[1].set_data_type(TensorProto::BOOL);
}
return out;
})
.SetDoc(R"DOC(
Convert sparse representations to dense with given indices.
Transforms a sparse representation of map<id, value> represented as `indices`
vector and `values` tensor into a compacted tensor where the first dimension
corresponds to each id provided in the mask argument. Missing values are filled
with the value of `default_value`. After running this op:
output[j, :] = values[i] // where mask[j] == indices[i]
output[j, ...] = default_value // when mask[j] doesn't appear in indices
If `lengths` is provided and not empty, an extra "batch" dimension is prepended
to the output.
`values` and `default_value` can have additional matching dimensions
(the operation is performed on the entire subtensor in this case).
For example, if `lengths` is supplied and `values` is a 1-D vector of floats
and `default_value` is a float scalar, the output is going to be a float
matrix of size `len(lengths) X len(mask)`.
)DOC")
.Arg(
"mask",
"list(int) argument with desired ids on the 'dense' output dimension")
.Arg(
"return_presence_mask",
"bool whether to return presence mask, false by default")
.Arg(
"max_skipped_indices",
"int argument representing the maximum number of invalid row ids that "
"can be skipped before returning an error. 50 by default")
.Input(0, "indices", "1-D int32/int64 tensor of concatenated ids of data")
.Input(1, "values", "Data tensor, first dimension has to match `indices`")
.Input(
2,
"default_value",
"Default value for the output if the id is not present in `indices`. "
"Must have the same type as `values` and the same shape, but without "
"the first dimension")
.Input(
3,
"lengths",
"Optional lengths to represent a batch of `indices` and `values`.")
.Output(
0,
"output",
"Output tensor of the same type as `values` of shape `[len(lengths), "
"len(mask)] + shape(default_value)` (if `lengths` is not provided the "
"first dimension is omitted)")
.Output(
1,
"presence_mask",
"Bool tensor of shape `[len(lengths), len(mask)]` (if `lengths` is not "
"provided the first dimension is omitted). True when a value for given "
"id was present, false otherwise.");
OPERATOR_SCHEMA(SparseToDenseMaskGradient)
.NumInputs(2, 3)
.NumOutputs(1)
.DisallowInputFillers() // TODO: enable the filler
.SetDoc(R"DOC(
The output is the gradient of the input value from SparseToDenseMask. The
gradient for default_value has not been implemented.
)DOC");
class GetSparseToDenseMaskGradient : public GradientMakerBase {
using GradientMakerBase::GradientMakerBase;
vector<OperatorDef> GetGradientDefs() override {
vector<string> blob_names{I(0), GO(0)};
// Add lengths blob if given
if (def_.input_size() == 4) {
blob_names.push_back(I(3));
}
return SingleGradientDef(
"SparseToDenseMaskGradient", "", blob_names, vector<string>{GI(1)});
}
};
REGISTER_GRADIENT(SparseToDenseMask, GetSparseToDenseMaskGradient);
} // namespace
} // namespace caffe2
// clang-format off
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
SparseToDenseMask,
"_caffe2::SparseToDenseMask("
"Tensor indices, "
"Tensor values, "
"Tensor default_value, "
"Tensor? lengths, "
"int[] mask, "
"bool? return_presence_mask = False, "
"int? max_skipped_indices = 50"
") -> (Tensor output, Tensor presence_mask)",
caffe2::SparseToDenseMaskOp<caffe2::CPUContext>);
// clang-format on