text: /**
   * Copyright (c) 2016-present, Facebook, Inc.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #ifndef CAFFE2_OPERATORS_RESHAPE_OP_H_
  #define CAFFE2_OPERATORS_RESHAPE_OP_H_
  
  #include "caffe2/core/common_omp.h"
  #include "caffe2/core/context.h"
  #include "caffe2/core/logging.h"
  #include "caffe2/core/operator.h"
  
  namespace caffe2 {
  
  // Takes a shape and data tensor and reshapes it
  template <typename F, class Context>
  class ReshapeOp : public Operator<Context> {
   public:
    USE_OPERATOR_CONTEXT_FUNCTIONS;
    ReshapeOp(const OperatorDef& operator_def, Workspace* ws)
        : Operator<Context>(operator_def, ws),
          new_shape_(OperatorBase::GetRepeatedArgument<int64_t>("shape")) {}
  
    bool RunOnDevice() override {
      if (InputSize() == 2) {
        return DispatchHelper<TensorTypes<int, int64_t>>::call(this, Input(1));
      }
      CAFFE_ENFORCE(
          OperatorBase::HasArgument("shape"), "Argument `shape` is missing.");
      return this->template DoRunWithType<int64_t>();
    }
  
    template <typename T>
    bool DoRunWithType() {
      auto& input = Input(0);
  
      vector<int64_t> actual_new_shape = new_shape_;
      if (InputSize() == 2) {
        CAFFE_ENFORCE(
            !OperatorBase::HasArgument("shape"),
            "New shape is specified by the input blob, do not pass in "
            "the argument `shape`.");
  
        auto& shape = Input(1);
        CAFFE_ENFORCE(shape.ndim() == 1, "Shape should be 1-D");
  
        const T* shape_data = shape.template data<T>();
  
        // Bit awkward, but needed so works on both CPU and CUDA contexts
        std::vector<T> tmpv(shape.size());
        context_.template CopyBytes<Context, CPUContext>(
            shape.size() * sizeof(T), shape_data, &tmpv[0]);
        actual_new_shape.assign(tmpv.begin(), tmpv.begin() + shape.size());
      }
  
      // Copy over the dimensions for those that are specified zero.
      for (int i = 0; i < actual_new_shape.size(); ++i) {
        if (actual_new_shape[i] == 0) {
          actual_new_shape[i] = input.dim(i);
        }
      }
  
      // Checks if the new shape is valid and fills in the missing dimension
      // specified by -1.
      // NOTE: At most one dimension can be -1.
      auto total_size = input.size_from_dim(0);
      T size = 1;
      int unknown_idx = -1;
      for (int i = 0; i < actual_new_shape.size(); ++i) {
        const auto dim = actual_new_shape[i];
        if (dim == -1) {
          CAFFE_ENFORCE(
              unknown_idx == -1,
              "Argument `shape` has more than one missing dimension.");
          unknown_idx = i;
        } else {
          size *= dim;
        }
      }
  
      if (unknown_idx != -1) {
        CAFFE_ENFORCE(
            total_size % size == 0,
            "Argument `shape` does not agree with the input data.",
            " (",
            total_size,
            " vs ",
            size,
            ")");
        actual_new_shape[unknown_idx] = total_size / size;
      } else {
        CAFFE_ENFORCE_EQ(
            total_size,
            size,
            "Argument `shape` does not agree with the input data.",
            " (",
            total_size,
            " != ",
            size,
            ")");
      }
  
      // Write the original shape to the second output.
      auto* old_shape = Output(1);
      old_shape->Resize(input.ndim());
      T* old_shape_data = old_shape->template mutable_data<T>();
      for (int i = 0; i < input.ndim(); ++i) {
        math::Set<T, Context>(1, input.dim(i), old_shape_data + i, &context_);
      }
  
      auto* output = Output(0);
      output->Resize(actual_new_shape);
      if (output != &input) {
        // If we are not doing in-place computation, a copy is needed.
        context_.template CopyBytes<Context, Context>(
            input.nbytes(),
            input.raw_data(),
            output->raw_mutable_data(input.meta()));
      }
  
      return true;
    }
  
   private:
    vector<int64_t> new_shape_;
  };
  
  } // namespace caffe2
  
  #endif // CAFFE2_OPERATORS_RESHAPE_OP_H_
  