pytorch  1.8.2 About: PyTorch provides Tensor computation (like NumPy) with strong GPU acceleration and Deep Neural Networks (in Python) built on a tape-based autograd system. LTS (Long Term Support) release.   Fossies Dox: pytorch-1.8.2.tar.gz  ("unofficial" and yet experimental doxygen-generated source code documentation)
batch_gather_ops.cc
Go to the documentation of this file.
2
3namespace caffe2 {
4
7
8OPERATOR_SCHEMA(BatchGather)
9 .NumInputs(2)
10 .NumOutputs(1)
11 .TensorInferenceFunction([](const OperatorDef& def,
15 const auto& data_dims = GetDimsVector(in[0]);
16 const auto& indices_dims = GetDimsVector(in[1]);
17
19 caffe2::gather_helper::calc_output_shape_vector<int>(
20 data_dims, indices_dims, 1, false);
22 return out;
23 })
24 .SetDoc(R"DOC(
25Batch gather operation, first dimension in DATA is the batch size.
26Given DATA tensor of rank r >= 2, and INDICES tensor of rank q >= 1, gather
27entries of the second outer dimension (axis == 1) of DATA indexed by INDICES,
28and concatenate them in an output tensor of rank q + (r - 1).
29
30Example:
31 DATA = [
32 [1.0, 1.2, 2.4, 4.5],
33 [2.3, 3.4, 3.6, 2.3],
34 [4.5, 5.7, 1.2, 4.5],
35 ]
36 INDICES = [0, 2]
37
38 OUTPUT = [
39 [1.0, 2.4],
40 [2.3, 3.6],
41 [4.5, 1.2],
42 ]
43)DOC")
44 .Input(0, "DATA", "Tensor of rank r >= 2.")
45 .Input(1, "INDICES", "Tensor of int32/int64 indices, of any rank q.")
46 .Output(0, "OUTPUT", "Tensor of rank q + (r - 1).")
47 .InheritOnnxSchema();
48
50
57 "",
58 vector<string>{I(Op::DATA), I(Op::INDICES), GO(0)},
59 vector<string>{GI(0)});
60 }
61};
62
64
65} // namespace caffe2
A helper class to index into arguments.
Definition: proto_utils.h:203
static vector< OperatorDef > SingleGradientDef(const Args &... args)
a helper function to allow one to create one single operator def, which is usually the case for many ...
Definition: blob.h:13
REGISTER_CPU_OPERATOR(ATen, ATenOp< CPUContext >)
parameter efficient embedding termed TT which can be plugged in into any model and trained end to end The benefits of our compressed TT layer are twofold instead of storing huge embedding it stores a sequence of much smaller dimensional and dimensional necessary for reconstructing the required which allows compressing the model significantly at the cost of a negligible performance drop the overall number of parameters can be relatively which allows to use larger batches or train efficiently in a case of limited resources DOC vector< int >
OPERATOR_SCHEMA(ATen)
return vector< TensorShape >
Definition: slice_op.cc:110
const auto & data_dims
ArgumentHelper helper(def)
const vector< TensorShape > & in
TensorShape CreateTensorShape(vector< T_I > dims, ::caffe2::TensorProto_DataType dt)
vector< int64_t > GetDimsVector(const TensorShape &shape)
vector< int > output_dims
true SparseLengthsFused8BitRowwiseOp< CPUContext, true >::WEIGHTS uint8 tensor obtained with INDICES
SparseLengths8BitsRowwiseOp< CPUContext, 0, 1 >::LENGTHS SetDoc(R"DOC( Variation of SparseLengthsMean operator, where DATA is stored using 8bits. DATA was quantized with 8Bit row-wise quantization (see doc to FloatToRowwiseQuantized8Bits operator). To restore DATA from 8Bit, we use additional input that stores scales and biases. )DOC") .Input(0
false SparseLengthsFused4BitRowwiseFakeFP16Op< CPUContext, false >::LENGTHS DATA