tensor.h

#ifndef SINGA_CORE_TENSOR_H_
#define SINGA_CORE_TENSOR_H_

#include <vector>
#include <tuple>
#include <memory>

#include "singa/core/common.h"
#include "singa/core/device.h"
#include "singa/proto/core.pb.h"
#include "singa/utils/logging.h"

using std::vector;
using std::tuple;
namespace singa {

typedef vector<size_t> Shape;
const size_t kDataWidth[] = {sizeof(float),  sizeof(float) / 2,
                             sizeof(int),    sizeof(char),
                             sizeof(double), sizeof(unsigned char)
                            };
inline size_t SizeOf(DataType t) {
  static_assert(kNumDataType == sizeof(kDataWidth) / sizeof(size_t),
                "Num of data types not match num of data width");
  CHECK_GT(kNumDataType, t);
  return kDataWidth[t];
}

class Tensor {
 public:
  ~Tensor();
  Tensor();

  explicit Tensor(const Shape &shape, DataType dtype = kFloat32);

  Tensor(const Shape &shape,
         std::shared_ptr<Device> dev,
         DataType dtype = kFloat32);

  Tensor(const Tensor &from);

  Tensor(Tensor &&from);

  // --------------------------------------------------------------------------
  // ---Following methods return info of the class without making any changes--
  // --------------------------------------------------------------------------

  Block *block() const { return block_; }

  std::shared_ptr<Device> device() const { return device_; }

  template <typename SType>
  const SType *data() const {
    return static_cast<const SType *>(block()->data());
  }

  const DataType data_type() const { return data_type_; }

  const Shape &shape() const { return shape_; }

  const size_t shape(const size_t idx) const {
    CHECK_LT(idx, shape_.size());
    return shape_.at(idx);
  }

  size_t nDim() const { return shape_.size(); }

  bool empty() const { return nDim() == 0; }

  bool transpose() const {
    if (!stride_.empty()) {
      auto last = stride_.front();
      for (auto s : stride_) {
        if (s > last && last > 0)  
          return true;
        if (s > 0)
          last = s;
      }
    }
    return false;
  }

  const vector<int>& stride() const { return stride_; }

  bool initailized() const {
    return block_ != nullptr && block_->initialized();
  }

  size_t Size() const {
    if (block_ == nullptr) return 0u;
    CHECK_EQ(block_->size() % SizeOf(data_type_), 0u);
    return block_->size() / SizeOf(data_type_);
  }

  size_t MemSize() const { return block_->size(); }

  template <typename SType>
  void GetValue(SType *value, const size_t num) {
    CHECK(device_ == defaultDevice);
    const SType* ptr = data<SType>();
    for (size_t i = 0; i < num; i++) value[i] = ptr[i];
  }

  void ToProto(singa::TensorProto *proto) const;

  float L1() const;

  float L2() const;
  // --------------------------------------------------------------------------
  // ---Following methods changes the internal data
  // --------------------------------------------------------------------------

  template <typename SType>
  void SetValue(const SType x);

  template <typename SType>
  void CopyDataFromHostPtr(const SType *src, const size_t num,
                           const size_t offset = 0);

  void CopyData(const Tensor &other);

  void FromProto(const singa::TensorProto &proto);


  void RepeatData(const vector<size_t>& repeats, int axis, int total_repeats,
                  const Tensor &other);

  // --------------------------------------------------------------------------
  // ---Following methods returns a new Tensor without change original tensor
  // --------------------------------------------------------------------------

  Tensor Repeat(const vector<size_t>& repeats, int axis,
                std::shared_ptr<Device> device = nullptr);

  Tensor Clone(std::shared_ptr<Device> device = nullptr) const;

  // --------------------------------------------------------------------------
  // ---Following methods change the tensor and return itself
  // --------------------------------------------------------------------------
  Tensor &operator=(const Tensor &in);

  Tensor &operator=(Tensor &&in);

  Tensor &operator+=(const Tensor &in);

  Tensor &operator-=(const Tensor &in);

  Tensor &operator*=(const Tensor &in);

  Tensor &operator/=(const Tensor &in);

  // Scalar operations.

  template <typename SType>
  Tensor &operator+=(const SType x);

  template <typename SType>
  Tensor &operator-=(const SType x);

  template <typename SType>
  Tensor &operator*=(const SType x);

  template <typename SType>
  Tensor &operator/=(const SType x);

  Tensor &Reshape(const Shape &shape);


  Tensor& Resize(const Shape& shape);

  Tensor& T();

  Tensor& Transpose();

  Tensor& Transpose(const vector<size_t> &axes);

  Tensor& Broadcast(const Shape& shape);

  Tensor& ResetLike(const Tensor &t);

  Tensor& AsType(const DataType type);

  Tensor& ToDevice(std::shared_ptr<Device> dev);

  Tensor& ToHost();

 protected:

  //generate strides automatically if stride field is not passed
  void generate_stride() {
    stride_.clear();
    if (shape_.size() == 0) {
      stride_.push_back(1);
      return;
    }

    size_t dim = Size();
    int cumulative_product = 1;
    for (size_t n = 0; n < shape_.size(); ++n) {
      cumulative_product = cumulative_product * shape_[n];
      stride_.push_back(dim / cumulative_product);
    }
  }

  void set_strides(const vector<int> new_strides) {
    stride_ = new_strides;
  }

 protected:
  DataType data_type_ = kFloat32;
  std::shared_ptr<Device> device_ = nullptr;
  Block *block_ = nullptr;
  Shape shape_ = {};
  vector<int> stride_ = {};
}; //end of tensor class


inline size_t Product(const Shape &shape, int start = 0, size_t len = 0) {
  if (len == 0) len = shape.size();
  if (len == 0) return 0;
  CHECK_LE(len, shape.size());
  size_t v = 1;
  for (unsigned int i = start; i < len; i++) v *= shape[i];
  return v;
}


inline void CheckDataTypeAndLang(const Tensor &in1, const Tensor &in2) {
  CHECK_EQ(in1.data_type(), in2.data_type());
  CHECK_EQ(in1.device()->lang(), in2.device()->lang());
}


template <typename FromType, typename ToType>
ToType TypeCast(const FromType &x) {
  // TODO(wangwei) cast fp16; prevent some casts, e.g., float to char
  return static_cast<ToType>(x);
}

Tensor Boradcast(const Shape& shape);

Tensor Reshape(const Tensor &in, const Shape &s);

Tensor Resize(const Tensor &in, const Shape &s);

Tensor Transpose(const Tensor& in);

Tensor Broadcast(const Tensor& in, const Shape& shape);

Tensor Transpose(const Tensor& in, const vector<size_t> &axes);

void CopyDataToFrom(Tensor *dst, const Tensor &src, const size_t num,
                    const size_t dst_offset = 0, const size_t src_offset = 0);


void RepeatDataToFrom(bool broadcast_flag, const vector<size_t>& repeats, int axis,
                      Tensor *dst, const Tensor &in, const size_t num);

// =============Element-wise operations====================================
Tensor Abs(const Tensor &in);
Tensor Exp(const Tensor &in);
Tensor Log(const Tensor &in);
Tensor ReLU(const Tensor &in);
Tensor Sigmoid(const Tensor &in);
Tensor Sign(const Tensor &in);
Tensor Sqrt(const Tensor &in);
Tensor Square(const Tensor &in);
Tensor Tanh(const Tensor &in);
Tensor Transform(const Tensor &in);

void Abs(const Tensor &in, Tensor *out);
void Exp(const Tensor &in, Tensor *out);
void Log(const Tensor &in, Tensor *out);
void ReLU(const Tensor &in, Tensor *out);
void Sigmoid(const Tensor &in, Tensor *out);
void Sign(const Tensor &in, Tensor *out);
void Sqrt(const Tensor &in, Tensor *out);
void Square(const Tensor &in, Tensor *out);
void Tanh(const Tensor &in, Tensor *out);
void Transform(const Tensor &in, Tensor *out);

template <typename SType>
Tensor Pow(const Tensor &in, const SType x);
template <typename SType>
void Pow(const Tensor &in, const SType x, Tensor *out);
Tensor Pow(const Tensor &base, const Tensor &exp);
void Pow(const Tensor &base, const Tensor &exp, Tensor *out);

template <typename SType>
Tensor operator<(const Tensor &in, const SType x);
template <typename SType>
void LT(const Tensor &in, const SType x, Tensor *out);

Tensor operator<(const Tensor &in1, const Tensor& in2);
void LT(const Tensor &in1, const Tensor& in2, Tensor *out);

template <typename SType>
Tensor operator<=(const Tensor &in, const SType x);
template <typename SType>
void LE(const Tensor &in, const SType x, Tensor *out);

Tensor operator<=(const Tensor &in1, const Tensor& in2);
void LE(const Tensor &in1, const Tensor& in2, Tensor *out);

template <typename SType>
Tensor operator>(const Tensor &in, const SType x);
template <typename SType>
void GT(const Tensor &in, const SType x, Tensor *out);

Tensor operator>(const Tensor &in1, const Tensor& in2);
void GT(const Tensor &in1, const Tensor& in2, Tensor *out);


template <typename SType>
Tensor operator>=(const Tensor &in, const SType x);
template <typename SType>
void GE(const Tensor &in, const SType x, Tensor *out);

Tensor operator>=(const Tensor &in1, const Tensor& in2);
void GE(const Tensor &in1, const Tensor& in2, Tensor *out);


Tensor operator+(const Tensor &lhs, const Tensor &rhs);
void Add(const Tensor &lhs, const Tensor &rhs, Tensor *out);
Tensor operator-(const Tensor &lhs, const Tensor &rhs);
void Sub(const Tensor &lhs, const Tensor &rhs, Tensor *out);
Tensor operator*(const Tensor &lhs, const Tensor &rhs);
void EltwiseMult(const Tensor &lhs, const Tensor &rhs, Tensor *out);
Tensor operator/(const Tensor &lhs, const Tensor &rhs);
void Div(const Tensor &lhs, const Tensor &rhs, Tensor *out);

template <typename SType>
Tensor operator+(const Tensor &in, const SType x);
template <typename SType>
void Add(const Tensor &in, const SType x, Tensor *out);

template <typename SType>
Tensor operator-(const Tensor &in, const SType x);
template <typename SType>
void Sub(const Tensor &in, const SType x, Tensor *out);

template <typename SType>
Tensor operator*(const Tensor &in, const SType x);
template <typename SType>
void EltwiseMult(const Tensor &in, const SType x, Tensor *out);

template <typename SType>
Tensor operator/(const Tensor &in, const SType x);
template <typename SType>
void Div(const Tensor &in, const SType x, Tensor *out);

template <typename SType>
Tensor Div(const SType x, const Tensor &in);
template <typename SType>
void Div(const SType x, const Tensor &in, Tensor *out);

template <typename SType = float>
SType Sum(const Tensor &in);


// ============Matrix (row/column) operations==================================
Tensor Average(const Tensor &in, const int axis);

void AddColumn(const Tensor &v, Tensor *M);
template <typename SType>
void AddColumn(const SType alpha, const SType beta, const Tensor &v,
               Tensor *out);
void AddRow(const Tensor &v, Tensor *out);
template <typename SType>
void AddRow(const SType alpha, const SType beta, const Tensor &v, Tensor *M);
void DivColumn(const Tensor &v, Tensor *M);
void DivRow(const Tensor &v, Tensor *M);
void MultColumn(const Tensor &v, Tensor *M);
void MultRow(const Tensor &v, Tensor *M);
Tensor SoftMax(const Tensor &in);

Tensor RowMax(const Tensor &in);
void SoftMax(const Tensor &in, Tensor *out);
void SubColumn(const Tensor &v, Tensor *M);
void SubRow(const Tensor &v, Tensor *M);
void SumColumns(const Tensor &M, Tensor *out);
void SumRows(const Tensor &M, Tensor *out);

Tensor Sum(const Tensor &in, const int axis);

// ================Random operations==========================================
template <typename SType>
void Bernoulli(const SType p, Tensor *out);
template <typename SType>
void Gaussian(const SType mean, const SType std, Tensor *out);
template <typename SType>
void Uniform(const SType low, const SType high, Tensor *out);

// ================Blas operations============================================
// TODO(wangwei) make amax/amin/asum a member function of tensor

template <typename SType>
void Axpy(SType alpha, const Tensor &in, Tensor *out);

Tensor Mult(const Tensor &A, const Tensor &B);
void Mult(const Tensor &A, const Tensor &B, Tensor *C);
template <typename SType>
void Mult(const SType alpha, const Tensor &A, const Tensor &B, const SType beta,
          Tensor *C);

// *****************
// Misc.
// ****************

void ComputeCrossEntropy(const Tensor &p, const Tensor &t, Tensor *loss);


void SoftmaxCrossEntropyBwd(const Tensor &t, Tensor *p);

Tensor CrossEntropyFwd(const Tensor& p, const Tensor& t);
Tensor SoftmaxCrossEntropyBwd(const Tensor& p, const Tensor& t);

Tensor CopyRows(const Tensor &in, const size_t start, const size_t end);
Tensor SliceRows(const Tensor &in, const size_t start, const size_t end);
Tensor SliceOn(const Tensor &in, const size_t start, const size_t end,
               int axis);
Tensor CopyColumns(const Tensor &in, const size_t start, const size_t end);
Tensor SliceColumns(const Tensor &in, const size_t start, const size_t end);
Tensor ConcatenateRows(const vector<Tensor> &in);
Tensor ConcatOn(const std::vector<Tensor> &in, int axis);
Tensor ConcatRows(const vector<Tensor> &in);
Tensor ConcatenateColumns(const vector<Tensor> &in);
Tensor ConcatColumns(const vector<Tensor> &in);
}  // namespace singa

#endif  // SINGA_CORE_TENSOR_H_