Caffe2 - C++ API
A deep learning, cross platform ML framework
gather_ranges_to_dense_op.cc
1 
17 #include "caffe2/operators/gather_ranges_to_dense_op.h"
18 
19 #include <cmath>
20 
21 namespace caffe2 {
22 namespace {
23 
24 OPERATOR_SCHEMA(GatherRangesToDense)
25  .NumInputs(2)
26  .NumOutputs(1, INT_MAX)
27  .SetDoc(R"DOC(
28 Given DATA tensor of rank 1, and RANGES tensor of rank 3, gather values
29 corresponding to each range into a separate output tensor.
30 
31 RANGES dimentions description:
32 1: represents list of examples within a batch
33 2: represents list features
34 3: two values which are start and length or a range (to be applied on DATA)
35 
36 Each feature has fixed lengths which are passed as lengths argument and a
37 separate tensor will be produced for each feature.
38 i.e. DATA.dim(1) = len(lengths) = NumOuptuts.
39 
40 Missing features (represented by empty ranges) filled with default_value.
41 
42 Example:
43  DATA = [1, 2, 3, 4, 5, 6, 7, 8]
44  RANGES = [
45  [
46  [2, 4],
47  [0, 2],
48  ],
49  [
50  [0, 0],
51  [6, 2],
52  ]
53  ]
54  lengths = [4, 2]
55  OUTPUT[0] = [[3, 4, 5, 6], [0, 0, 0, 0]]
56  OUTPUT[1] = [[1, 2], [7, 8]]
57 )DOC")
58  .Input(0, "DATA", "Tensor of rank 1.")
59  .Input(
60  1,
61  "RANGES",
62  "Tensor of int32/int64 ranges, of dims (N, M, 2). "
63  "Where N is number of examples and M is a size of each example. "
64  "Last dimention represents a range in the format (start, lengths)")
65  .Output(0, "OUTPUT", "1-D tensor of size sum of range lengths")
66  .Arg("lengths", "Expected lengths for ranges")
67  .TensorInferenceFunction([](const OperatorDef& def,
68  const vector<TensorShape>& in) {
69  ArgumentHelper helper(def);
70  auto lengths = helper.GetRepeatedArgument<int>("lengths");
71  CAFFE_ENFORCE_EQ(in[0].dims_size(), 1, "DATA should be 1-D tensor.");
72  CAFFE_ENFORCE_EQ(in[1].dims_size(), 3, "RANGES should be 3-D tensor.");
73  CAFFE_ENFORCE_GT(lengths.size(), 0, "lengths should be non-empty.");
74  std::vector<TensorShape> out(lengths.size());
75  for (int i = 0; i < lengths.size(); ++i) {
76  out[i].set_data_type(in[0].data_type());
77  out[i].add_dims(in[1].dims(0));
78  out[i].add_dims(lengths[i]);
79  }
80  return out;
81  });
82 
83 REGISTER_CPU_OPERATOR(GatherRangesToDense, GatherRangesToDenseOp<CPUContext>);
84 NO_GRADIENT(GatherRangesToDense);
85 
86 } // namespace
87 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.