|
#pragma once
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/operator.h"
|
|
namespace caffe2 {
|
|
// RecordShapeOp records the shape of the input tensor to a vector of int. You
|
// mostly don't need this operator explicitly, and it is mostly used in the
|
// autodiff process.
|
template <class Context>
|
class ShapeOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit ShapeOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
axes_(OperatorBase ::GetRepeatedArgument<int>("axes")) {}
|
|
bool RunOnDevice() override {
|
auto& data = Input(DATA);
|
|
int numDims = data.dim();
|
int numAxes = axes_.size();
|
if (numAxes == 0) {
|
auto* output = Output(0, {numDims}, at::dtype<int64_t>());
|
int64_t* output_data = output->template mutable_data<int64_t>();
|
context_.CopyBytesSameDevice(
|
numDims * sizeof(int64_t), data.sizes().data(), output_data);
|
return true;
|
}
|
|
auto* output = Output(0, {numAxes}, at::dtype<int64_t>());
|
auto src = reinterpret_cast<const char*>(data.sizes().data());
|
auto out = reinterpret_cast<char*>(output->template mutable_data<int64_t>());
|
for (int i = 0; i < numAxes; i++) {
|
auto axis = axes_[i];
|
CAFFE_ENFORCE_LT(axis, numDims, "Axis out of range");
|
CAFFE_ENFORCE_GE(axis, 0, "Each axis should be non-negative");
|
context_.CopyBytesSameDevice(
|
sizeof(int64_t), src + axis * sizeof(int64_t), out);
|
out += sizeof(int64_t);
|
}
|
return true;
|
}
|
|
INPUT_TAGS(DATA);
|
|
private:
|
vector<int> axes_;
|
};
|
|
} // namespace caffe2
|
