caffe2/operators/concat_split_op.h - platform/external/pytorch - Git at Google

 #ifndef CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_
 #define CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_

 #include "caffe2/core/context.h"
 #include "caffe2/core/operator.h"
 #include "caffe2/core/types.h"
 #include "caffe2/utils/math.h"
 #include "caffe2/utils/string_utils.h"
 #include <c10/util/accumulate.h>
 #include <c10/util/irange.h>

 namespace caffe2 {

 template <class Context>
 class SplitOp final : public Operator<Context> {
  public:
   static const int kSplitOpInputSize = 2;

   USE_OPERATOR_CONTEXT_FUNCTIONS;
   template <class... Args>
   explicit SplitOp(Args&&... args)
       : Operator<Context>(std::forward<Args>(args)...),
         split_(this->template GetRepeatedArgument<int>("split")) {
     CAFFE_ENFORCE(
         !(OperatorBase::HasArgument("axis") &&
           OperatorBase::HasArgument("order")),
         "You shouldn't specify both the dim to split, and the order "
         "in the case of 4-D images.");
     if (OperatorBase::HasArgument("axis")) {
       axis_ = this->template GetSingleArgument<int>("axis", -1);
       // only exists for computing the gradient of a Concat with 'add_axis'
       add_axis_ = this->template GetSingleArgument<int>("add_axis", 0);
     } else {
       axis_ = GetDimFromOrderString(
           this->template GetSingleArgument<string>("order", "NCHW"));
       add_axis_ = 0;
     }
   }

   bool RunOnDevice() override;

  protected:
   int axis_;
   int add_axis_;
   vector<int> split_;
   // Input: X, optionally split
   // The split tensor is stored in CPU.
 };

 template <class Context>
 class SplitByLengthsOp final : public Operator<Context> {
  public:
   USE_OPERATOR_CONTEXT_FUNCTIONS;
   template <class... Args>
   explicit SplitByLengthsOp(Args&&... args)
       : Operator<Context>(std::forward<Args>(args)...) {
     CAFFE_ENFORCE(
         !(OperatorBase::HasArgument("axis") &&
           OperatorBase::HasArgument("order")),
         "You shouldn't specify both the dim to split, and the order "
         "in the case of 4-D images.");
     if (OperatorBase::HasArgument("axis")) {
       axis_ = this->template GetSingleArgument<int>("axis", 0);
     } else {
       axis_ = GetDimFromOrderString(
           this->template GetSingleArgument<string>("order", "NCHW"));
     }
     scaling_ =
         this->template GetSingleArgument<bool>("use_scaling_lengths", false);
   }

   bool RunOnDevice() override;

  protected:
   int axis_;
   bool scaling_;
   Tensor inclusive_scan_buffer_{Context::GetDeviceType()};
   Tensor inclusive_scan_length_buffer_{Context::GetDeviceType()};
   // Input: X, optionally split
   // The split tensor is stored in CPU.
   Tensor lengths_host_{CPU};
 };

 template <class Context>
 class ConcatOp final : public Operator<Context> {
  public:
   USE_OPERATOR_CONTEXT_FUNCTIONS;
   template <class... Args>
   explicit ConcatOp(Args&&... args)
       : Operator<Context>(std::forward<Args>(args)...) {
     CAFFE_ENFORCE(
         !(OperatorBase::HasArgument("axis") &&
           OperatorBase::HasArgument("order")),
         "You shouldn't specify both the dim to concat, and the order "
         "in the case of 4-D images.");
     if (OperatorBase::HasArgument("axis")) {
       axis_ = this->template GetSingleArgument<int>("axis", -1);
       add_axis_ = this->template GetSingleArgument<int>("add_axis", 0);
     } else {
       axis_ = GetDimFromOrderString(
           this->template GetSingleArgument<string>("order", "NCHW"));
       add_axis_ = 0;
     }
   }

   bool RunOnDevice() override;

  protected:
   int axis_;
   int add_axis_;
   // Input: a number of tensors. Output: Y, split
   // The split are stored in CPU.
 };

 // Implementations
 template <class Context>
 bool SplitOp<Context>::RunOnDevice() {
   auto& input = Input(0);
   int canonical_axis = input.canonical_axis_index(axis_);
   CAFFE_ENFORCE_LT(
       canonical_axis, input.dim(), "Axis not in input ndim range.");
   const int input_channels = input.dim32(canonical_axis);
   const int* axis_data;
   vector<int> equal_split;
   if (InputSize() == kSplitOpInputSize) {
     // We obtain split from the input tensor.
     CAFFE_ENFORCE_EQ(
         split_.size(),
         0,
         "If you set split with an input blob, do not pass in "
         "split in the argument.");
     auto& split_tensor = this->template Input<Tensor>(1, CPU);
     CAFFE_ENFORCE_EQ(split_tensor.numel(), OutputSize());
     axis_data = split_tensor.template data<int>();
   } else if (split_.size() == 0) {
     CAFFE_ENFORCE_EQ(
         input_channels % OutputSize(),
         0,
         "If you did not specify split explicitly, the number of "
         "input channels:",
         input_channels,
         " should be divisible by the output size:",
         OutputSize(),
         ".");
     equal_split.resize(OutputSize(), input_channels / OutputSize());
     axis_data = equal_split.data();
   } else {
     // We obtain split from the parameters.
     CAFFE_ENFORCE_EQ(
         split_.size(),
         OutputSize(),
         "The number of splits specified should be equal to the "
         "number of outputs.");
     axis_data = split_.data();
   }

   CAFFE_ENFORCE_EQ(
       add_axis_ ? OutputSize()
                 : std::accumulate(axis_data, axis_data + OutputSize(), 0),
       input_channels,
       "Sum of split dimensions do not match: should be ",
       input_channels);
   vector<int64_t> output_dims(input.sizes().vec());
   int before = 1, after = 1;
   for (const auto i : c10::irange(canonical_axis)) {
     before *= input.dim32(i);
   }
   for (int i = canonical_axis + 1; i < input.dim(); ++i) {
     after *= input.dim32(i);
   }
   if (add_axis_) {
     output_dims.erase(output_dims.begin() + canonical_axis);
   }

   const auto *const input_ptr = static_cast<const char*>(input.raw_data());

   size_t input_offset = 0;
   for (const auto i : c10::irange(OutputSize())) {
     auto *const output = Output(i);
     const auto axis_dim = add_axis_ ? 1 : axis_data[i];
     if (!add_axis_) {
       output_dims[canonical_axis] = axis_data[i];
     }
     output->Resize(output_dims);

     // We need `output_ptr` before the early exit since
     // `raw_mutable_data` sets the output's data type
     auto *const output_ptr = output->raw_mutable_data(input.dtype());

     if (input_ptr == nullptr || output_ptr == nullptr) {
       continue;
     }

     math::CopyMatrix<Context>(
         input.itemsize(),
         before,
         axis_dim * after,
         input_ptr + input_offset,
         input.dim32(canonical_axis) * after,
         output_ptr,
         axis_dim * after,
         &context_,
         input.dtype().copy());
     input_offset += axis_dim * after * input.itemsize();
   }
   return true;
 }

 // Implementations
 template <class Context>
 bool SplitByLengthsOp<Context>::RunOnDevice() {
   auto& input = Input(0);
   auto lengths_length = Input(1).dim(0);
   int32_t* length_data;

   if (this->InputIsTensorType(1, CPU)) {
     length_data = Input(1).template data<int32_t>();
   } else {
     // Length input in CUDA context
     auto& input_length = Input(1);
     lengths_host_ = TensorCPU(input_length, CPU);
     length_data = lengths_host_.template data<int32_t>();
   }

   CAFFE_ENFORCE_EQ(
       lengths_length % OutputSize(),
       0,
       "len(Lengths) ",
       lengths_length,
       "should be divisible by OutputSize() ",
       OutputSize(),
       ".");
   int canonical_axis = input.canonical_axis_index(axis_);
   CAFFE_ENFORCE_LT(
       canonical_axis, input.dim(), "Axis not in input ndim range.");
   const int input_channels = input.dim32(canonical_axis);
   const auto* axis_data = length_data;

   auto sum_lengths = std::accumulate(axis_data, axis_data + lengths_length, 0);

   if (scaling_) {
     CAFFE_ENFORCE_EQ(
         input_channels % (sum_lengths ? sum_lengths : 1),
         0,
         "Input channels ",
         input_channels,
         " should be divisible by ",
         sum_lengths);
   } else {
     CAFFE_ENFORCE_EQ(
         sum_lengths,
         input_channels,
         "Input channels should be equal to split dimensions sum, ",
         input_channels,
         " vs ",
         sum_lengths);
   }
   vector<int64_t> output_dims(input.sizes().vec());
   int before = input.size_to_dim(canonical_axis);
   int after = input.size_from_dim(canonical_axis + 1);
   size_t input_offset = 0;
   auto dim_multiplier = sum_lengths ? (input_channels / sum_lengths) : 1;

   if (!scaling_) {
     dim_multiplier = 1;
   }

   for (const auto i : c10::irange(OutputSize())) {
     auto* output = Output(i);
     const auto* axis_offset = axis_data + lengths_length / OutputSize() * i;
     auto axis_dim =
         dim_multiplier *
         std::accumulate(
             axis_offset, axis_offset + lengths_length / OutputSize(), 0);
     output_dims[canonical_axis] = axis_dim;
     output->Resize(output_dims);
     math::CopyMatrix<Context>(
         input.itemsize(),
         before,
         axis_dim * after,
         static_cast<const char*>(input.raw_data()) + input_offset,
         input.dim32(canonical_axis) * after,
         output->raw_mutable_data(input.dtype()),
         axis_dim * after,
         &context_,
         input.dtype().copy());
     input_offset += axis_dim * after * input.itemsize();
   }
   return true;
 }

 template <class Context>
 bool ConcatOp<Context>::RunOnDevice() {
   auto *const output = Output(0);

   // We can override default options(Context::GetDeviceType())
   // by explicitly passing in device type we want
   Tensor *const split = Output(
       1, at::IntArrayRef({InputSize()}), at::dtype<int>().device(CPU));
   int *const axis_data = split->template mutable_data<int>();
   auto& input_zero = Input(0);
   int adj_size = input_zero.dim() + (add_axis_ ? 1 : 0);
   int canonical_axis = canonical_axis_index_(axis_, adj_size);
   CAFFE_ENFORCE_LT(canonical_axis, adj_size, "Axis not in input ndim range.");
   for (const auto i : c10::irange(1, InputSize())) {
     CAFFE_ENFORCE_EQ(
         Input(i).dtype(),
         input_zero.dtype(),
         "All inputs must have the same type, expected: ",
         input_zero.dtype().name(),
         " but got: ",
         Input(i).dtype().name(),
         " for input: ",
         i);
   }

   int before = 1, after = 1;
   vector<int64_t> output_dims(input_zero.sizes().vec());
   for (const auto i : c10::irange(input_zero.dim())) {
     if (i == canonical_axis && !add_axis_) {
       continue;
     }
     int dim = input_zero.dim32(i);
     if (i < canonical_axis) {
       before *= dim;
     } else { // i > canonical_axis || i == canonical_axis && add_axis_
       after *= dim;
     }
     // check the input dims are compatible.
     for (const auto j : c10::irange(1, InputSize())) {
       int dim_j = Input(j).dim32(i);
       CAFFE_ENFORCE_EQ(
           dim,
           dim_j,
           "Expect dimension = ",
           dim,
           " got ",
           dim_j,
           " at axis = ",
           i,
           " for input: ",
           j,
           ". The input tensors can only have different dimensions "
           "when arg 'add_axis' = 0 and along the axis = ",
           canonical_axis,
           " <",
           Input(0).sizes(),
           "> vs <",
           Input(j).sizes(),
           ">.");
     }
   }

   int output_channels = 0;
   for (const auto i : c10::irange(InputSize())) {
     axis_data[i] = add_axis_ ? 1 : Input(i).dim32(canonical_axis);
     output_channels += axis_data[i];
   }
   if (add_axis_) {
     output_dims.insert(output_dims.begin() + canonical_axis, output_channels);
   } else {
     output_dims[canonical_axis] = output_channels;
   }

   output->Resize(output_dims);
   auto *const output_ptr = static_cast<char*>(output->raw_mutable_data(input_zero.dtype()));
   if(output_ptr == nullptr){
     return true;
   }

   size_t output_offset = 0;
   for (const auto i : c10::irange(InputSize())) {
     auto& input = Input(i);
     auto axis_dim = add_axis_ ? 1 : input.dim32(canonical_axis);
     math::CopyMatrix<Context>(
         input.itemsize(),
         before,
         axis_dim * after,
         input.raw_data(),
         axis_dim * after,
         output_ptr + output_offset,
         output_channels * after,
         &context_,
         input_zero.dtype().copy());
     output_offset += axis_dim * after * input.itemsize();
   }
   return true;
 }

 OpSchema::Cost CostInferenceForConcat(
     const OperatorDef& def,
     const std::vector<TensorShape>& in);

 std::vector<TensorShape> TensorInferenceForConcat(
     const OperatorDef& def,
     const std::vector<TensorShape>& in);

 } // namespace caffe2

 #endif // CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_
	#ifndef CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_
	#define CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_

	#include "caffe2/core/context.h"
	#include "caffe2/core/operator.h"
	#include "caffe2/core/types.h"
	#include "caffe2/utils/math.h"
	#include "caffe2/utils/string_utils.h"
	#include <c10/util/accumulate.h>
	#include <c10/util/irange.h>

	namespace caffe2 {

	template <class Context>
	class SplitOp final : public Operator<Context> {
	public:
	static const int kSplitOpInputSize = 2;

	USE_OPERATOR_CONTEXT_FUNCTIONS;
	template <class... Args>
	explicit SplitOp(Args&&... args)
	: Operator<Context>(std::forward<Args>(args)...),
	split_(this->template GetRepeatedArgument<int>("split")) {
	CAFFE_ENFORCE(
	!(OperatorBase::HasArgument("axis") &&
	OperatorBase::HasArgument("order")),
	"You shouldn't specify both the dim to split, and the order "
	"in the case of 4-D images.");
	if (OperatorBase::HasArgument("axis")) {
	axis_ = this->template GetSingleArgument<int>("axis", -1);
	// only exists for computing the gradient of a Concat with 'add_axis'
	add_axis_ = this->template GetSingleArgument<int>("add_axis", 0);
	} else {
	axis_ = GetDimFromOrderString(
	this->template GetSingleArgument<string>("order", "NCHW"));
	add_axis_ = 0;
	}
	}

	bool RunOnDevice() override;

	protected:
	int axis_;
	int add_axis_;
	vector<int> split_;
	// Input: X, optionally split
	// The split tensor is stored in CPU.
	};

	template <class Context>
	class SplitByLengthsOp final : public Operator<Context> {
	public:
	USE_OPERATOR_CONTEXT_FUNCTIONS;
	template <class... Args>
	explicit SplitByLengthsOp(Args&&... args)
	: Operator<Context>(std::forward<Args>(args)...) {
	CAFFE_ENFORCE(
	!(OperatorBase::HasArgument("axis") &&
	OperatorBase::HasArgument("order")),
	"You shouldn't specify both the dim to split, and the order "
	"in the case of 4-D images.");
	if (OperatorBase::HasArgument("axis")) {
	axis_ = this->template GetSingleArgument<int>("axis", 0);
	} else {
	axis_ = GetDimFromOrderString(
	this->template GetSingleArgument<string>("order", "NCHW"));
	}
	scaling_ =
	this->template GetSingleArgument<bool>("use_scaling_lengths", false);
	}

	bool RunOnDevice() override;

	protected:
	int axis_;
	bool scaling_;
	Tensor inclusive_scan_buffer_{Context::GetDeviceType()};
	Tensor inclusive_scan_length_buffer_{Context::GetDeviceType()};
	// Input: X, optionally split
	// The split tensor is stored in CPU.
	Tensor lengths_host_{CPU};
	};

	template <class Context>
	class ConcatOp final : public Operator<Context> {
	public:
	USE_OPERATOR_CONTEXT_FUNCTIONS;
	template <class... Args>
	explicit ConcatOp(Args&&... args)
	: Operator<Context>(std::forward<Args>(args)...) {
	CAFFE_ENFORCE(
	!(OperatorBase::HasArgument("axis") &&
	OperatorBase::HasArgument("order")),
	"You shouldn't specify both the dim to concat, and the order "
	"in the case of 4-D images.");
	if (OperatorBase::HasArgument("axis")) {
	axis_ = this->template GetSingleArgument<int>("axis", -1);
	add_axis_ = this->template GetSingleArgument<int>("add_axis", 0);
	} else {
	axis_ = GetDimFromOrderString(
	this->template GetSingleArgument<string>("order", "NCHW"));
	add_axis_ = 0;
	}
	}

	bool RunOnDevice() override;

	protected:
	int axis_;
	int add_axis_;
	// Input: a number of tensors. Output: Y, split
	// The split are stored in CPU.
	};

	// Implementations
	template <class Context>
	bool SplitOp<Context>::RunOnDevice() {
	auto& input = Input(0);
	int canonical_axis = input.canonical_axis_index(axis_);
	CAFFE_ENFORCE_LT(
	canonical_axis, input.dim(), "Axis not in input ndim range.");
	const int input_channels = input.dim32(canonical_axis);
	const int* axis_data;
	vector<int> equal_split;
	if (InputSize() == kSplitOpInputSize) {
	// We obtain split from the input tensor.
	CAFFE_ENFORCE_EQ(
	split_.size(),
	0,
	"If you set split with an input blob, do not pass in "
	"split in the argument.");
	auto& split_tensor = this->template Input<Tensor>(1, CPU);
	CAFFE_ENFORCE_EQ(split_tensor.numel(), OutputSize());
	axis_data = split_tensor.template data<int>();
	} else if (split_.size() == 0) {
	CAFFE_ENFORCE_EQ(
	input_channels % OutputSize(),
	0,
	"If you did not specify split explicitly, the number of "
	"input channels:",
	input_channels,
	" should be divisible by the output size:",
	OutputSize(),
	".");
	equal_split.resize(OutputSize(), input_channels / OutputSize());
	axis_data = equal_split.data();
	} else {
	// We obtain split from the parameters.
	CAFFE_ENFORCE_EQ(
	split_.size(),
	OutputSize(),
	"The number of splits specified should be equal to the "
	"number of outputs.");
	axis_data = split_.data();
	}

	CAFFE_ENFORCE_EQ(
	add_axis_ ? OutputSize()
	: std::accumulate(axis_data, axis_data + OutputSize(), 0),
	input_channels,
	"Sum of split dimensions do not match: should be ",
	input_channels);
	vector<int64_t> output_dims(input.sizes().vec());
	int before = 1, after = 1;
	for (const auto i : c10::irange(canonical_axis)) {
	before *= input.dim32(i);
	}
	for (int i = canonical_axis + 1; i < input.dim(); ++i) {
	after *= input.dim32(i);
	}
	if (add_axis_) {
	output_dims.erase(output_dims.begin() + canonical_axis);
	}

	const auto const input_ptr = static_cast<const char>(input.raw_data());

	size_t input_offset = 0;
	for (const auto i : c10::irange(OutputSize())) {
	auto *const output = Output(i);
	const auto axis_dim = add_axis_ ? 1 : axis_data[i];
	if (!add_axis_) {
	output_dims[canonical_axis] = axis_data[i];
	}
	output->Resize(output_dims);

	// We need `output_ptr` before the early exit since
	// `raw_mutable_data` sets the output's data type
	auto *const output_ptr = output->raw_mutable_data(input.dtype());

	if (input_ptr == nullptr \|\| output_ptr == nullptr) {
	continue;
	}

	math::CopyMatrix<Context>(
	input.itemsize(),
	before,
	axis_dim * after,
	input_ptr + input_offset,
	input.dim32(canonical_axis) * after,
	output_ptr,
	axis_dim * after,
	&context_,
	input.dtype().copy());
	input_offset += axis_dim * after * input.itemsize();
	}
	return true;
	}

	// Implementations
	template <class Context>
	bool SplitByLengthsOp<Context>::RunOnDevice() {
	auto& input = Input(0);
	auto lengths_length = Input(1).dim(0);
	int32_t* length_data;

	if (this->InputIsTensorType(1, CPU)) {
	length_data = Input(1).template data<int32_t>();
	} else {
	// Length input in CUDA context
	auto& input_length = Input(1);
	lengths_host_ = TensorCPU(input_length, CPU);
	length_data = lengths_host_.template data<int32_t>();
	}

	CAFFE_ENFORCE_EQ(
	lengths_length % OutputSize(),
	0,
	"len(Lengths) ",
	lengths_length,
	"should be divisible by OutputSize() ",
	OutputSize(),
	".");
	int canonical_axis = input.canonical_axis_index(axis_);
	CAFFE_ENFORCE_LT(
	canonical_axis, input.dim(), "Axis not in input ndim range.");
	const int input_channels = input.dim32(canonical_axis);
	const auto* axis_data = length_data;

	auto sum_lengths = std::accumulate(axis_data, axis_data + lengths_length, 0);

	if (scaling_) {
	CAFFE_ENFORCE_EQ(
	input_channels % (sum_lengths ? sum_lengths : 1),
	0,
	"Input channels ",
	input_channels,
	" should be divisible by ",
	sum_lengths);
	} else {
	CAFFE_ENFORCE_EQ(
	sum_lengths,
	input_channels,
	"Input channels should be equal to split dimensions sum, ",
	input_channels,
	" vs ",
	sum_lengths);
	}
	vector<int64_t> output_dims(input.sizes().vec());
	int before = input.size_to_dim(canonical_axis);
	int after = input.size_from_dim(canonical_axis + 1);
	size_t input_offset = 0;
	auto dim_multiplier = sum_lengths ? (input_channels / sum_lengths) : 1;

	if (!scaling_) {
	dim_multiplier = 1;
	}

	for (const auto i : c10::irange(OutputSize())) {
	auto* output = Output(i);
	const auto* axis_offset = axis_data + lengths_length / OutputSize() * i;
	auto axis_dim =
	dim_multiplier *
	std::accumulate(
	axis_offset, axis_offset + lengths_length / OutputSize(), 0);
	output_dims[canonical_axis] = axis_dim;
	output->Resize(output_dims);
	math::CopyMatrix<Context>(
	input.itemsize(),
	before,
	axis_dim * after,
	static_cast<const char*>(input.raw_data()) + input_offset,
	input.dim32(canonical_axis) * after,
	output->raw_mutable_data(input.dtype()),
	axis_dim * after,
	&context_,
	input.dtype().copy());
	input_offset += axis_dim * after * input.itemsize();
	}
	return true;
	}

	template <class Context>
	bool ConcatOp<Context>::RunOnDevice() {
	auto *const output = Output(0);

	// We can override default options(Context::GetDeviceType())
	// by explicitly passing in device type we want
	Tensor *const split = Output(
	1, at::IntArrayRef({InputSize()}), at::dtype<int>().device(CPU));
	int *const axis_data = split->template mutable_data<int>();
	auto& input_zero = Input(0);
	int adj_size = input_zero.dim() + (add_axis_ ? 1 : 0);
	int canonical_axis = canonical_axis_index_(axis_, adj_size);
	CAFFE_ENFORCE_LT(canonical_axis, adj_size, "Axis not in input ndim range.");
	for (const auto i : c10::irange(1, InputSize())) {
	CAFFE_ENFORCE_EQ(
	Input(i).dtype(),
	input_zero.dtype(),
	"All inputs must have the same type, expected: ",
	input_zero.dtype().name(),
	" but got: ",
	Input(i).dtype().name(),
	" for input: ",
	i);
	}

	int before = 1, after = 1;
	vector<int64_t> output_dims(input_zero.sizes().vec());
	for (const auto i : c10::irange(input_zero.dim())) {
	if (i == canonical_axis && !add_axis_) {
	continue;
	}
	int dim = input_zero.dim32(i);
	if (i < canonical_axis) {
	before *= dim;
	} else { // i > canonical_axis \|\| i == canonical_axis && add_axis_
	after *= dim;
	}
	// check the input dims are compatible.
	for (const auto j : c10::irange(1, InputSize())) {
	int dim_j = Input(j).dim32(i);
	CAFFE_ENFORCE_EQ(
	dim,
	dim_j,
	"Expect dimension = ",
	dim,
	" got ",
	dim_j,
	" at axis = ",
	i,
	" for input: ",
	j,
	". The input tensors can only have different dimensions "
	"when arg 'add_axis' = 0 and along the axis = ",
	canonical_axis,
	" <",
	Input(0).sizes(),
	"> vs <",
	Input(j).sizes(),
	">.");
	}
	}

	int output_channels = 0;
	for (const auto i : c10::irange(InputSize())) {
	axis_data[i] = add_axis_ ? 1 : Input(i).dim32(canonical_axis);
	output_channels += axis_data[i];
	}
	if (add_axis_) {
	output_dims.insert(output_dims.begin() + canonical_axis, output_channels);
	} else {
	output_dims[canonical_axis] = output_channels;
	}

	output->Resize(output_dims);
	auto const output_ptr = static_cast<char>(output->raw_mutable_data(input_zero.dtype()));
	if(output_ptr == nullptr){
	return true;
	}

	size_t output_offset = 0;
	for (const auto i : c10::irange(InputSize())) {
	auto& input = Input(i);
	auto axis_dim = add_axis_ ? 1 : input.dim32(canonical_axis);
	math::CopyMatrix<Context>(
	input.itemsize(),
	before,
	axis_dim * after,
	input.raw_data(),
	axis_dim * after,
	output_ptr + output_offset,
	output_channels * after,
	&context_,
	input_zero.dtype().copy());
	output_offset += axis_dim * after * input.itemsize();
	}
	return true;
	}

	OpSchema::Cost CostInferenceForConcat(
	const OperatorDef& def,
	const std::vector<TensorShape>& in);

	std::vector<TensorShape> TensorInferenceForConcat(
	const OperatorDef& def,
	const std::vector<TensorShape>& in);

	} // namespace caffe2

	#endif // CAFFE2_OPERATORS_CONCAT_SPLIT_OP_H_