text: #ifndef CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_
  #define CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_
  
  #include "caffe2/core/context.h"
  #include "caffe2/core/export_caffe2_op_to_c10.h"
  #include "caffe2/core/operator.h"
  
  C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(PiecewiseLinearTransform);
  
  namespace caffe2 {
  
  template <typename T, class Context>
  class PiecewiseLinearTransformOp final : public Operator<Context> {
   public:
    USE_OPERATOR_CONTEXT_FUNCTIONS;
  
    template <class... Args>
    explicit PiecewiseLinearTransformOp(Args&&... args)
        : Operator<Context>(std::forward<Args>(args)...) {
      binary_ = this->template GetSingleArgument<bool>("binary", false);
  
      // Retrieve transform params (i.e., the linear functions).
      bounds_from_arg_ = this->template GetRepeatedArgument<T>("bounds");
      slopes_from_arg_ = this->template GetRepeatedArgument<T>("slopes");
      intercepts_from_arg_ = this->template GetRepeatedArgument<T>("intercepts");
      transform_param_from_arg_ = CheckTransParamFromArg();
    }
  
    bool RunOnDevice() override {
      return binary_ ? TransformBinary() : TransformGeneral();
    }
  
   private:
    // num_func_per_group is the number of pieces of linear functions of
    // each group.
    // num_group: The number of groups of linear functions. Each group is for
    // transforming one column of predictions.
    void InferNumFunctionsPerGroup(
        const int64_t num_bounds,
        const int64_t num_slopes,
        const int64_t num_intercepts,
        int64_t* num_func_per_group,
        int64_t* num_group) {
      CAFFE_ENFORCE_EQ(num_slopes, num_intercepts);
  
      // This is based on the facts:
      // 1. in each group, the num of bounds minus the num of slopes is 1;
      // 2. each group has the same number of pieces.
      *num_group = num_bounds - num_slopes;
      CAFFE_ENFORCE_GT(*num_group, 0);
      if (binary_) {
        CAFFE_ENFORCE_EQ(*num_group, 1);
      }
      *num_func_per_group = num_slopes / *num_group;
      CAFFE_ENFORCE_GT(*num_func_per_group, 0);
      CAFFE_ENFORCE_EQ(num_slopes % *num_group, 0);
    }
  
    bool CheckBoundsSorted(
        const T* bounds,
        const int64_t num_bounds_per_group,
        const int64_t num_group) {
      const T* start = bounds;
      for (int64_t i = 0; i < num_group; i++) {
        if (!std::is_sorted(start, start + num_bounds_per_group)) {
          return false;
        }
        start += num_bounds_per_group;
      }
      return true;
    }
  
    // Returns true if the transform params from arg are valid.
    // Otherwise, we will assume the transform params will pass from Input blobs.
    bool CheckTransParamFromArg() {
      int good_param = 0;
      good_param += bounds_from_arg_.size() > 0;
      good_param += slopes_from_arg_.size() > 0;
      good_param += intercepts_from_arg_.size() > 0;
      CAFFE_ENFORCE(
          good_param == 0 || good_param == 3,
          "bounds, slopes, intercepts must be all set or all not set");
      if (good_param == 3) {
        int64_t num_func_per_group;
        int64_t num_group;
        InferNumFunctionsPerGroup(
            bounds_from_arg_.size(),
            slopes_from_arg_.size(),
            intercepts_from_arg_.size(),
            &num_func_per_group,
            &num_group);
        CAFFE_ENFORCE(
            CheckBoundsSorted(
                bounds_from_arg_.data(), num_func_per_group + 1, num_group),
            "bounds must be sorted for each group");
      }
  
      return good_param == 3;
    }
  
    void setUpTensors(int64_t& num_func_per_group, int64_t& num_group, int64_t M);
  
    void GetTransParamData(
        const T** bounds,
        const T** slopes,
        const T** intercepts,
        int64_t* num_func_per_group,
        int64_t* num_group) {
      int64_t num_bounds;
      int64_t num_slopes;
      int64_t num_intercepts;
  
      if (transform_param_from_arg_) {
        CAFFE_ENFORCE_EQ(InputSize(), 1);
        *bounds = bounds_from_arg_.data();
        *slopes = slopes_from_arg_.data();
        *intercepts = intercepts_from_arg_.data();
        num_bounds = bounds_from_arg_.size();
        num_slopes = slopes_from_arg_.size();
        num_intercepts = intercepts_from_arg_.size();
      } else {
        CAFFE_ENFORCE_EQ(InputSize(), 4);
        auto& bounds_input = Input(BOUNDS);
        auto& slopes_input = Input(SLOPES);
        auto& intercepts_input = Input(INTERCEPTS);
        *bounds = bounds_input.template data<T>();
        *slopes = slopes_input.template data<T>();
        *intercepts = intercepts_input.template data<T>();
        num_bounds = bounds_input.numel();
        num_slopes = slopes_input.numel();
        num_intercepts = intercepts_input.numel();
      }
      InferNumFunctionsPerGroup(
          num_bounds, num_slopes, num_intercepts, num_func_per_group, num_group);
    }
  
    bool TransformGeneral() {
      auto& X = Input(0);
  
      CAFFE_ENFORCE_EQ(X.dim(), 2);
      int64_t N = X.dim32(0);
      int64_t M = X.dim32(1);
      auto* Y = Output(0, X.sizes(), at::dtype<T>());
      const auto* Xdata = X.template data<T>();
      T* Ydata = Y->template mutable_data<T>();
  
      const T* bounds;
      const T* slopes;
      const T* intercepts;
      int64_t num_func_per_group;
      int64_t num_group;
      GetTransParamData(
          &bounds, &slopes, &intercepts, &num_func_per_group, &num_group);
      CAFFE_ENFORCE_EQ(num_group, M);
  
      for (int64_t j = 0; j < M; ++j) {
        const T* bounds_group = bounds + j * (num_func_per_group + 1);
        const T* slopes_group = slopes + j * num_func_per_group;
        const T* intercepts_group = intercepts + j * num_func_per_group;
        for (int64_t i = 0; i < N; ++i) {
          Ydata[i * M + j] = PiecewiseLinearTransform(
              Xdata[i * M + j],
              bounds_group,
              slopes_group,
              intercepts_group,
              num_func_per_group);
        }
      }
      return true;
    }
  
    bool TransformBinary() {
      auto& X = Input(PREDICTIONS);
  
      CAFFE_ENFORCE(X.dim() == 1 || X.dim() == 2);
      int64_t N = X.dim32(0);
      int64_t M = X.dim() == 2 ? X.dim32(1) : 1;
      CAFFE_ENFORCE(
          M == 1 || M == 2,
          "If binary is set to true, the input must be Nx2 or Nx1 tensor");
      auto* Y = Output(0, X.sizes(), at::dtype<T>());
      const auto* Xdata = X.template data<T>();
      T* Ydata = Y->template mutable_data<T>();
  
      const T* bounds;
      const T* slopes;
      const T* intercepts;
      int64_t num_func_per_group;
      int64_t num_group;
      GetTransParamData(
          &bounds, &slopes, &intercepts, &num_func_per_group, &num_group);
      CAFFE_ENFORCE_EQ(num_group, 1);
  
      if (M == 1) {
        for (int64_t i = 0; i < N; ++i) {
          Ydata[i] = PiecewiseLinearTransform(
              Xdata[i], bounds, slopes, intercepts, num_func_per_group);
        }
      } else {
        for (int64_t i = 0; i < N; ++i) {
          Ydata[i * M + 1] = PiecewiseLinearTransform(
              Xdata[i * M + 1], bounds, slopes, intercepts, num_func_per_group);
          Ydata[i * M] = 1.0f - Ydata[i * M + 1];
        }
      }
  
      return true;
    }
  
    T PiecewiseLinearTransform(
        const T x,
        const T* bounds,
        const T* slopes,
        const T* intercepts,
        const int64_t num_func_per_group) {
      T y = 0;
      // deal with samples out of bounds
      // make it the same as the upper/lower bound value
      if (x <= bounds[0]) {
        y = slopes[0] * bounds[0] + intercepts[0];
      } else if (x >= bounds[num_func_per_group]) {
        y = slopes[num_func_per_group - 1] * bounds[num_func_per_group] +
            intercepts[num_func_per_group - 1];
      } else {
        auto low_bound =
            std::lower_bound(bounds, bounds + num_func_per_group + 1, x);
        int bounds_idx = low_bound - bounds - 1;
        // compute the piecewise linear transformation as Y
        y = slopes[bounds_idx] * x + intercepts[bounds_idx];
      }
      return y;
    }
  
   private:
    bool binary_;
    vector<T> bounds_from_arg_;
    vector<T> slopes_from_arg_;
    vector<T> intercepts_from_arg_;
  
    Tensor bounds_device_{Context::GetDeviceType()};
    Tensor intercepts_device_{Context::GetDeviceType()};
    Tensor slopes_device_{Context::GetDeviceType()};
    bool gpu_copied_ = false;
  
    // If true, the piecewise linear functions are passed through args,
    // otherwise, they are passed through Input blobs.
    bool transform_param_from_arg_;
  
    INPUT_TAGS(PREDICTIONS, BOUNDS, SLOPES, INTERCEPTS);
  };
  
  } // namespace caffe2
  
  #endif // CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_
  