text: #include "caffe2/operators/sequence_ops.h"
  #include "caffe2/core/operator.h"
  #include "caffe2/core/tensor.h"
  
  namespace caffe2 {
  
  vector<TensorShape> TensorInferenceForAddPadding(
      const OperatorDef& def,
      const vector<TensorShape>& in) {
    ArgumentHelper helper(def);
    const int padding_width = helper.GetSingleArgument<int>("padding_width", 1);
    const int end_padding_width = helper.GetSingleArgument<int>("end_padding_width", padding_width);
    CAFFE_ENFORCE_GT(in.size(), 0);
    CAFFE_ENFORCE_GE(in[0].dims_size(), 1);
    if (in.size() > 1) {
      CAFFE_ENFORCE_EQ(in[1].dims_size(), 1);
    }
  
    const auto num_paddings = (in.size() == 1 ? 1 : in[1].dims(0));
    vector<int> out_shape(in[0].dims().begin(), in[0].dims().end());
    out_shape[0] += (padding_width + end_padding_width) * num_paddings;
  
    if (def.output_size() == 1) {
      return vector<TensorShape>{CreateTensorShape(out_shape, in[0].data_type())};
    } else {
      return vector<TensorShape>{
        CreateTensorShape(out_shape, in[0].data_type()),
        CreateTensorShape(vector<int>(1, num_paddings), TensorProto::INT32)};
    }
  }
  
  template <>
  template <typename T>
  void GatherPaddingOp<CPUContext>::GatherPadding(
      const int outer_size,
      const int lengths_size,
      const int block_size,
      const int pad_width,
      const T* in_ptr,
      const int* lengths_ptr,
      T* padding_start_ptr,
      T* padding_end_ptr) {
    CAFFE_ENFORCE(
        (!std::is_same<bool, T>::value),
        "GatherPadding should not be executed on an input of type bool, as "
        "addition is not properly defined with booleans.");
    int64_t total_length = 0;
    for (int i = 0; i < lengths_size; ++i) {
      // check total length consistency
      const auto length = lengths_ptr[i];
      total_length += length;
      CAFFE_ENFORCE_LE(total_length, outer_size);
      // accumulate start paddings
      for (int j = 0; j < startPaddingWidth_; ++j) {
        for (int k = 0; k < block_size; ++k) {
          // Note: MSVC warns about unsafe use of type bool in operation.
          // This is now guarded by a CAFFE_ENFORCE so we can suppress it.
          #pragma warning(suppress: 4804)
          padding_start_ptr[k] += in_ptr[k];
        }
        in_ptr += block_size;
      }
      in_ptr += block_size * (length - pad_width);
      // accumulate end paddings
      for (int j = 0; j < endPaddingWidth_; ++j) {
        for (int k = 0; k < block_size; ++k) {
          #pragma warning(suppress: 4804)
          padding_end_ptr[k] += in_ptr[k];
        }
        in_ptr += block_size;
      }
    }
  }
  
  template <>
  template <typename T>
  bool RemovePaddingOp<CPUContext>::DoRunWithType() {
    const auto& in = Input(0);
    CAFFE_ENFORCE_GE(in.dim(), 1);
    const int32_t outer_size = in.sizes()[0];
    const auto block_size = std::accumulate(
        // NOLINTNEXTLINE(modernize-use-transparent-functors)
        in.sizes().begin() + 1, in.sizes().end(), 1, std::multiplies<int64_t>());
    const auto pad_width = startPaddingWidth_ + endPaddingWidth_;
  
    // if no lengths is provided, assume it is a single full-span entry
    const int32_t* lengths_ptr = &outer_size;
    int64_t lengths_size = 1;
    if (InputSize() > 1) {
      const auto& lengths = Input(1);
      lengths_ptr = lengths.data<int32_t>();
      lengths_size = lengths.numel();
    }
  
    auto out_dims = in.sizes().vec();
    out_dims[0] -= pad_width * lengths_size;
    auto* out = Output(0, std::move(out_dims), at::dtype<T>());
  
    const auto* in_ptr = in.template data<T>();
    auto* out_ptr = out->template mutable_data<T>();
    int64_t total_length = 0;
    for (int i = 0; i < lengths_size; ++i) {
      // check that total length is consistent
      const auto length = lengths_ptr[i];
      total_length += length;
      CAFFE_ENFORCE_LE(total_length, outer_size);
      std::copy(
          in_ptr + block_size * startPaddingWidth_,
          in_ptr + block_size * (length - endPaddingWidth_),
          out_ptr);
      in_ptr += block_size * length;
      out_ptr += block_size * (length - pad_width);
    }
    if (OutputSize() == 1) {
      return true;
    }
  
    auto* lengths_out = Output(1, {lengths_size}, at::dtype<int32_t>());
    std::transform(
        lengths_ptr,
        lengths_ptr + lengths_size,
        lengths_out->template mutable_data<int32_t>(),
        [pad_width](int32_t x) { return x - pad_width; });
    return true;
  }
  
  template <>
  template <typename T>
  bool AddPaddingOp<CPUContext>::MakePadding(
      const T* in_ptr,
      T* out_ptr,
      const int32_t* lengths_ptr,
      int32_t lengths_size,
      int32_t outer_size,
      const T* padding_start_ptr,
      const T* padding_end_ptr,
      int64_t block_size) {
    if (!lengths_ptr) {
      lengths_ptr = &outer_size;
    }
  
    int64_t total_length = 0;
    for (int i = 0; i < lengths_size; ++i) {
      // check that total length is consistent
      const auto length = lengths_ptr[i];
      total_length += length;
      CAFFE_ENFORCE_LE(total_length, outer_size);
      // copy padding before
      if (!padding_start_ptr) {
        memset(out_ptr, 0, block_size * startPaddingWidth_ * sizeof(T));
        out_ptr += block_size * startPaddingWidth_;
      } else {
        for (int j = 0; j < startPaddingWidth_; ++j) {
          std::copy(padding_start_ptr, padding_start_ptr + block_size, out_ptr);
          out_ptr += block_size;
        }
      }
      // copy payload
      const auto num_elems = block_size * length;
      std::copy(in_ptr, in_ptr + num_elems, out_ptr);
      in_ptr += num_elems;
      out_ptr += num_elems;
      // copy padding after
      if (!padding_end_ptr) {
        memset(out_ptr, 0, block_size * endPaddingWidth_ * sizeof(T));
        out_ptr += block_size * endPaddingWidth_;
      } else {
        for (int j = 0; j < endPaddingWidth_; ++j) {
          std::copy(padding_end_ptr, padding_end_ptr + block_size, out_ptr);
          out_ptr += block_size;
        }
      }
    }
    if (OutputSize() == 1) {
      return true;
    }
  
    auto* lengths_out = Output(1, {lengths_size}, at::dtype<int32_t>());
    const auto pad_width = startPaddingWidth_ + endPaddingWidth_;
    std::transform(
        lengths_ptr,
        lengths_ptr + lengths_size,
        lengths_out->template mutable_data<int32_t>(),
        [pad_width](int32_t x) { return x + pad_width; });
    return true;
  }
  
  template <>
  bool PadEmptySamplesOp<CPUContext>::RunOnDevice() {
    auto& lengths = Input(0);
    auto* lengthsPtr = lengths.template data<int32_t>();
    CAFFE_ENFORCE(lengths.dim() == 1, "LENGTH should be 1-D");
    CAFFE_ENFORCE(InputSize() >= 1, "Input size must be no less than 1");
  
    int needPadding = 0;
    int sumLen = 0;
    for (int i = 0; i < lengths.numel(); ++i) {
      if (lengthsPtr[i] == 0) {
        needPadding++;
      }
      sumLen += lengthsPtr[i];
    }
  
    auto* out_lengths = Output(0, {lengths.numel()}, at::dtype<int32_t>());
    auto* outLengthsPtr = out_lengths->template mutable_data<int32_t>();
    for (int i = 0; i < lengths.numel(); ++i) {
      if (lengthsPtr[i] == 0) {
        outLengthsPtr[i] = 1;
      } else {
        outLengthsPtr[i] = lengthsPtr[i];
      }
    }
  
    for (int k = 0; k < InputSize() - 1; k++) {
      auto& features = Input(1 + k);
      CAFFE_ENFORCE(features.dim() >= 1, "FEATURE should at least 1-D");
      CAFFE_ENFORCE(
          features.size(0) == sumLen, "FEATURE and LENGTH should be consistent");
      const auto block_size = features.size_from_dim(1);
  
      auto* out_features = Output(1 + k);
      auto outDim = features.sizes().vec();
      outDim.at(0) += needPadding;
      out_features->Resize(outDim);
      auto dst =
          static_cast<char*>(out_features->raw_mutable_data(features.dtype()));
      auto src_base = static_cast<const char*>(features.raw_data());
      // copy data and add padding index as zero
      Tensor zero{CPU};
      zero.Resize(block_size);
      auto zeroPtr = static_cast<char*>(zero.raw_mutable_data(features.dtype()));
      // TODO Handle other composite types, such as vector<...>
      if (!features.dtype().Match<std::string>()) {
        memset(zeroPtr, 0, zero.nbytes());
      }
      int start_dest = 0;
      int start_src = 0;
      for (int i = 0; i < lengths.numel(); ++i) {
        if (lengthsPtr[i] == 0) {
          context_.CopyItemsSameDevice(
              features.dtype(),
              block_size,
              zeroPtr,
              dst + start_dest * features.dtype().itemsize());
          start_dest += block_size;
        } else {
          auto src = src_base + start_src * features.dtype().itemsize();
          context_.CopyItemsSameDevice(
              features.dtype(),
              lengthsPtr[i] * block_size,
              src,
              dst + start_dest * features.dtype().itemsize());
          // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
          start_src += lengthsPtr[i] * block_size;
          // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
          start_dest += lengthsPtr[i] * block_size;
        }
      }
    }
    return true;
  }
  
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_CPU_OPERATOR(AddPadding, AddPaddingOp<CPUContext>);
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_CPU_OPERATOR(RemovePadding, RemovePaddingOp<CPUContext>);
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_CPU_OPERATOR(GatherPadding, GatherPaddingOp<CPUContext>);
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_CPU_OPERATOR(PadEmptySamples, PadEmptySamplesOp<CPUContext>);
  
  struct GetAddPaddingGradient : public GradientMakerBase {
    using GradientMakerBase::GradientMakerBase;
    vector<OperatorDef> GetGradientDefs() override {
      // whether to provide lengths as input to gradient
      vector<std::string> g_inputs{GO(0)};
      if (Def().input_size() > 1) {
        CAFFE_ENFORCE(Def().output_size() > 1);
        g_inputs.push_back(O(1));
      }
  
      vector<OperatorDef> ops;
      // gradient on the data
      ops.push_back(CreateOperatorDef(
          "RemovePadding", "", g_inputs, vector<string>{GI(0)}));
      // gradient on the start_padding (and end_padding)
      if (Def().input_size() >= 3) {
        std::vector<string> padding_grads{GI(2)};
        if (Def().input_size() == 4) {
          padding_grads.push_back(GI(3));
        }
        // NOLINTNEXTLINE(performance-unnecessary-copy-initialization)
        auto g_inputs2 = g_inputs;
        ops.push_back(
            CreateOperatorDef("GatherPadding", "", g_inputs2, padding_grads));
      }
      return ops;
    }
  };
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_GRADIENT(AddPadding, GetAddPaddingGradient);
  
  struct GetRemovePaddingGradient : public GradientMakerBase {
    using GradientMakerBase::GradientMakerBase;
    vector<OperatorDef> GetGradientDefs() override {
      // whether to provide lengths as input to gradient
      vector<std::string> g_inputs{GO(0)};
      if (Def().input_size() > 1) {
        CAFFE_ENFORCE(Def().output_size() > 1);
        g_inputs.push_back(O(1));
      }
  
      return SingleGradientDef("AddPadding", "", g_inputs, vector<string>{GI(0)});
    }
  };
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  REGISTER_GRADIENT(RemovePadding, GetRemovePaddingGradient);
  
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  OPERATOR_SCHEMA(AddPadding)
      .NumInputs(1, 4)
      .NumOutputs(1, 2)
      .TensorInferenceFunction(
          OpSchema::NeedsAllInputShapes(TensorInferenceForAddPadding))
      .SetDoc(R"DOC(
  Given a partitioned tensor $T<N, D_1, ..., D_n>$, where the partitions are
  defined as ranges on its outer-most (slowest varying) dimension $N$,
  return a tensor $T<(N + 2 * padding\_width), D_1, ..., D_n>$ with paddings
  added to the start and end of each range.
  
  Optionally, different paddings can be provided for beginning and end.
  Paddings provided must be a tensor $T<D_1, ..., D_n>$. If no padding is
  provided, add zero padding. If no lengths vector is provided, add padding
  only once, at the start and end of data.
  
  Github Links:
  
  - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/sequence_ops.cc
  
  <details>
  
  <summary> <b>Example</b> </summary>
  
  **Code**
  
  ```
  workspace.ResetWorkspace()
  
  op = core.CreateOperator(
      "AddPadding",
      ["X", "lengths"],
      ["Y", "lengths_out"],
      padding_width=1
  
  )
  
  workspace.FeedBlob("X", (np.random.rand(3,2,2).astype(np.float32)))
  workspace.FeedBlob("lengths", np.array([3]).astype(np.int32))
  
  print("X:", workspace.FetchBlob("X"))
  workspace.RunOperatorOnce(op)
  print("Y:", workspace.FetchBlob("Y"))
  print("lengths_out:", workspace.FetchBlob("lengths_out"))
  ```
  
  **Result**
  
  ```
  X: [[[0.2531572  0.4588472 ]
    [0.45140603 0.61161053]]
  
   [[0.92500854 0.8045306 ]
    [0.03356671 0.30233648]]
  
   [[0.4660227  0.6287745 ]
    [0.79372746 0.08609265]]]
  Y: [[[0.         0.        ]
    [0.         0.        ]]
  
   [[0.2531572  0.4588472 ]
    [0.45140603 0.61161053]]
  
   [[0.92500854 0.8045306 ]
    [0.03356671 0.30233648]]
  
   [[0.4660227  0.6287745 ]
    [0.79372746 0.08609265]]
  
   [[0.         0.        ]
    [0.         0.        ]]]
  lengths_out: [5]
  ```
  
  </details>
  
  )DOC")
      .Arg(
          "padding_width",
          "*(type: int)* Number of copies of padding to add around each range.")
      .Arg(
          "end_padding_width",
          "*(type: int)* [OPTIONAL] Specifies a different end-padding width. If "
          "this is not set, will use same as `padding_width`.")
      .Input(
          0,
          "data_in",
          "*(type: Tensor)* Input data ($T<N, D_1, ..., D_n>$).")
      .Input(
          1,
          "lengths",
          "*(type: Tensor`<int>`)* Number of elements in each range. "
          "sum(lengths) = N.")
      .Input(
          2,
          "start_padding",
          "*(type: Tensor`<int>`)* [OPTIONAL] Padding data for range start "
          "($T<D_1, ..., D_n>$).")
      .Input(
          3,
          "end_padding",
          "*(type: Tensor`<int>`)* [OPTIONAL] Padding for range end. If not "
          "provided, `start_padding` is used ($T<D_1, ..., D_n>$).")
      .Output(
          0,
          "data_out",
          "*(type: Tensor)* Padded data tensor ($T<N + 2*padding_width, "
          "D_1, ..., D_n>$).")
      .Output(
          1,
          "lengths_out",
          "*(type: Tensor`<int>`)* [OPTIONAL] Lengths for each padded range.");
  
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  OPERATOR_SCHEMA(RemovePadding)
      .NumInputs(1, 2)
      .NumOutputs(1, 2)
      .SetDoc(R"DOC(
  Remove padding around the edges of each segment of the input data. This is the
  reverse operation of **AddPadding**, and uses the same arguments and conventions
  for input and output data format.
  
  Github Links:
  
  - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/sequence_ops.cc
  
  <details>
  
  <summary> <b>Example</b> </summary>
  
  **Code**
  
  ```
  workspace.ResetWorkspace()
  
  addpad_op = core.CreateOperator(
      "AddPadding",
      ["X", "lengths_add"],
      ["Y", "lengths_out_add"],
      padding_width=1
  )
  
  rmpad_op = core.CreateOperator(
      "RemovePadding",
      ["Y", "lengths_rm"],
      ["Z", "lengths_out_rm"],
      padding_width=1
  )
  
  workspace.FeedBlob("X", (np.random.randint(20, size=(3,5))))
  workspace.FeedBlob("lengths_add", np.array([3]).astype(np.int32))
  workspace.FeedBlob("lengths_rm", np.array([5]).astype(np.int32))
  
  print("X:", workspace.FetchBlob("X"))
  workspace.RunOperatorOnce(addpad_op)
  print("Y:", workspace.FetchBlob("Y"))
  print("lengths_out_add:", workspace.FetchBlob("lengths_out_add"))
  
  workspace.RunOperatorOnce(rmpad_op)
  print("Z:", workspace.FetchBlob("Z"))
  print("lengths_out_rm:", workspace.FetchBlob("lengths_out_rm"))
  ```
  
  **Result**
  
  ```
  X: [[17 19  1  9  1]
   [19  3  5 19  1]
   [16  0  0  0  4]]
  Y: [[ 0  0  0  0  0]
   [17 19  1  9  1]
   [19  3  5 19  1]
   [16  0  0  0  4]
   [ 0  0  0  0  0]]
  lengths_out_add: [5]
  Z: [[17 19  1  9  1]
   [19  3  5 19  1]
   [16  0  0  0  4]]
  lengths_out_rm: [3]
  ```
  
  </details>
  
  )DOC")
      .Arg(
          "padding_width",
          "*(type: int)* Outer-size of padding to remove around each range.")
      .Arg(
          "end_padding_width",
          "*(type: int)* [OPTIONAL] Specifies a different end-padding width. "
          "If this is not set, will use same as `padding_width`.")
      .Input(
          0,
          "data_in",
          "Input tensor ($T<N, D_1, ..., D_n>$).")
      .Input(
          1,
          "lengths",
          "*(type: Tensor`<int>`)* Number of elements in each range. "
          "sum(lengths) = N. If not provided, considers all data as a single "
          "segment.")
      .Output(
          0,
          "data_out",
          "*(type: Tensor)* Padded data tensor "
          "($T<N + 2*padding_width, D_1, ..., D_n>$).")
      .Output(
          1,
          "lengths_out",
          "*(type: Tensor`<int>`)* [OPTIONAL] Lengths for each padded range.");
  
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  OPERATOR_SCHEMA(GatherPadding)
      .NumInputs(2)
      .NumOutputs(1, 2)
      .SetDoc(R"DOC(
  Gather the sum of start and end paddings in a padded input sequence. Used in
  order to compute the gradients of AddPadding w.r.t the padding tensors.
  )DOC")
      .Arg("padding_width", "Outer-size of padding present around each range.")
      .Arg(
          "end_padding_width",
          "(Optional) Specifies a different end-padding width.")
      .Input(0, "data_in", "T<N, D1..., Dn> Padded input data")
      .Input(
          1,
          "lengths",
          "(i64) Num of elements in each range. sum(lengths) = N. "
          "If not provided, considers all data as a single segment.")
      .Output(
          0,
          "padding_sum",
          "Sum of all start paddings, or of all "
          "paddings if end_padding_sum is not provided.")
      .Output(
          1,
          "end_padding_sum",
          "T<D1..., Dn> Sum of all end paddings, if provided.");
  
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
  OPERATOR_SCHEMA(PadEmptySamples)
      .NumInputs(1, INT_MAX)
      .NumOutputs(1, INT_MAX)
      .SetDoc(R"DOC(
  Pad empty field given lengths and index features,
  
  Input(0) is a blob pointing to the lengths of samples in one batch,
  [Input(1),... Input(num_fields)] a list of tensors containing the data for
  each field of the features.
  
  PadEmptySamples is thread safe.
  )DOC")
      .Input(0, "lengths", "A blob containing a pointer to the lengths.")
      .Output(
          0,
          "out_lengths",
          "Tensor containing lengths with empty sample padded.");
  
  } // namespace caffe2
  