stats_put_ops.h

#include <limits>
#include "caffe2/core/operator.h"
#include "caffe2/core/stats.h"
#include "caffe2/core/tensor.h"
#include "caffe2/core/types.h"

namespace caffe2 {

template <typename T>
struct TemplatePutOp : public Operator<CPUContext> {
  explicit TemplatePutOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator(operator_def, ws),
        given_name_(GetSingleArgument<std::string>(
            "stat_name",
            operator_def.input().Get(0))),
        magnitude_expand_(GetSingleArgument<int64_t>("magnitude_expand", 1)),
        bound_(GetSingleArgument<bool>("bound", false)),
        has_default_(HasSingleArgumentOfType<float>("default_value")),
        default_value_(GetSingleArgument<float>("default_value", 0.0)),
        stat_(given_name_) {}

  bool RunOnDevice() override {
    return DispatchHelper<TensorTypes<
        int,
        float,
        uint8_t,
        int8_t,
        uint16_t,
        int16_t,
        int64_t,
        at::Half,
        double>>::call(this, Input(0));
  }

  template <typename V>
  bool DoRunWithType() {
    V input = default_value_;

    // If we receive an empty tensor
    if (Input(0).template data<V>()) {
      input = *Input(0).template data<V>();
    } else if (!has_default_) {
      CAFFE_THROW(
          "Default value must be provided when recieving empty tensors for ",
          given_name_);
    }

    int64_t bound_value =
        std::numeric_limits<int64_t>::max() / magnitude_expand_;

    int64_t int_value;
    if (bound_) {
      if (isNan(input)) {
        int_value = 0;
      } else if (input <= -bound_value) {
        int_value = std::numeric_limits<int64_t>::min();
      } else if (input >= bound_value) {
        int_value = std::numeric_limits<int64_t>::max();
      } else {
        int_value = input * magnitude_expand_;
      }
    } else {
      CAFFE_ENFORCE(
          std::abs(static_cast<int64_t>(input)) < bound_value,
          "Input value is too large for the given magnitude expansion!");
      CAFFE_ENFORCE(!isNan(input), "Input value cannot be NaN!");
      int_value = input * magnitude_expand_;
    }

    CAFFE_EVENT(stat_, stat_value, int_value);

    return true;
  }

 private:
  const std::string given_name_;
  const int64_t magnitude_expand_;
  const bool bound_;
  const bool has_default_;
  const float default_value_;
  T stat_;

  template <typename V>
  bool isNan(V input) {
    /*
    Checks if given number of is NaN, while being permissive with different
    implementations of the standard libraries between operating systems.

    Uses the preperties of NaN, defined by IEEE.
    https://www.gnu.org/software/libc/manual/html_node/Infinity-and-NaN.html
    */
    return input != input;
  }
};
} // namespace caffe2