tuple.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 MXNET_TUPLE_H_
#define MXNET_TUPLE_H_

#include <vector>
#include <type_traits>
#include <algorithm>
#include <utility>
#include <iostream>
#include <string>
#include "nnvm/op_attr_types.h"
#include "nnvm/graph_attr_types.h"
#include "nnvm/graph.h"
#include "nnvm/pass.h"

namespace mxnet {

template<typename ValueType>
class Tuple {
 public:
  Tuple() = default;
  inline ~Tuple() {
    delete [] data_heap_;
  }
  inline Tuple(const Tuple<ValueType>& s) {
    if (s.ndim() == -1) {
      this->SetDim(-1);
    } else {
      this->assign(s.begin(), s.end());
    }
  }
  inline Tuple(std::initializer_list<ValueType> init) {
    this->assign(init.begin(), init.end());
  }
  inline Tuple(std::vector<ValueType> init) {  // NOLINT(runtime/explicit)
    this->assign(init.begin(), init.end());
  }
  inline Tuple(Tuple<ValueType>&& src) {   // NOLINT(runtime/explicit)
    this->swap(src);
  }
  template<typename RandomAccessIterator>
  inline Tuple(RandomAccessIterator begin,
               RandomAccessIterator end) {
    this->assign(begin, end);
  }
  template<typename RandomAccessIterator>
  inline void assign(RandomAccessIterator begin,
                     RandomAccessIterator end) {
    this->SetDim(end - begin);
    CHECK_GE(ndim(), 0);
    std::copy(begin, end, this->begin());
  }
  inline void swap(Tuple<ValueType>& other) {  // NOLINT(*)
    std::swap(ndim_, other.ndim_);
    std::swap(num_heap_allocated_, other.num_heap_allocated_);
    std::swap(data_stack_, other.data_stack_);
    std::swap(data_heap_, other.data_heap_);
  }
  inline Tuple<ValueType>& operator=(const Tuple<ValueType>& src) {
    if (src.ndim() == -1) {
      this->SetDim(-1);
    } else {
      this->assign(src.begin(), src.end());
    }
    return *this;
  }
  inline Tuple<ValueType>& operator=(Tuple<ValueType>&& src) {
    Tuple<ValueType>(std::move(src)).swap(*this);
    return *this;
  }
  inline Tuple<ValueType> &operator=(std::initializer_list<ValueType> init) {
    this->assign(init.begin(), init.end());
    return *this;
  }
  inline bool operator==(const Tuple<ValueType> &s) const {
    if (ndim_ != s.ndim_) return false;
    if (ndim() == -1) return true;
    return std::equal(begin(), end(), s.begin());
  }
  inline bool operator!=(const Tuple<ValueType> &s) const {
    return !(*this == s);
  }
  inline const ValueType *begin() const {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  inline ValueType *begin() {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  inline const ValueType* end() const {
    return ndim_ <= kStackCache ? (data_stack_ + ndim_): (data_heap_ + ndim_);
  }
  inline ValueType* end() {
    return ndim_ <= kStackCache ? (data_stack_ + ndim_): (data_heap_ + ndim_);
  }
  inline int ndim() const {
    return ndim_;
  }
  inline ValueType& operator[](int i) {
    CHECK(i >= 0 && i < ndim()) << "index = " << i << " must be in range [0, " << ndim() << ")";
    return begin()[i];
  }
  inline const ValueType& operator[](int i) const {
    CHECK(i >= 0 && i < ndim()) << "index = " << i << " must be in range [0, " << ndim() << ")";
    return begin()[i];
  }
  inline void Save(dmlc::JSONWriter* writer) const {
    std::vector<ValueType> tmp(begin(), end());
    writer->Write(tmp);
  }
  inline void Load(dmlc::JSONReader* reader) {
    std::vector<ValueType> tmp;
    reader->Read(&tmp);
    this->assign(tmp.begin(), tmp.end());
  }
  friend std::ostream &operator<<(std::ostream &os, const Tuple<ValueType> &t) {
    if (t.ndim() == -1) {
      // If t is an unknown shape, return string "None".
      // This is consistent with returning unknown shape in Python and generating
      // C++ operator APIs by OpWrapperGenerator.py (defaultString) in cpp-package.
      os << "None";
      return os;
    }
    os << '[';
    const ValueType* begin = t.begin();
    const ValueType* end = t.end();
    for (const ValueType* it = begin; it != end; ++it) {
      if (it != begin) os << ',';
      os << *it;
    }
    os << ']';
    return os;
  }
  friend std::istream &operator>>(std::istream &is, Tuple<ValueType> &t) {
    // get (
    while (true) {
      char ch = is.peek();
      if (isdigit(ch) || ch == '-') {
        ValueType idx;
        if (is >> idx) {
          t.assign(&idx, &idx + 1);
        }
        return is;
      }
      is.get();
      if (ch == '(' || ch == '[') break;
      if (!isspace(ch)) {
        is.setstate(std::ios::failbit);
        return is;
      }
    }
    // Handle empty tuple. A tensor whose shape is an empty tuple
    // represents a scalar with ndim = 0.
    while (isspace(is.peek())) {
      is.get();
    }
    if (is.peek() == ')' || is.peek() == ']') {
      is.get();
      t.SetDim(0);
      return is;
    }
    // Handle non-empty tuple
    ValueType idx;
    std::vector<ValueType> tmp;
    while (is >> idx) {
      tmp.push_back(idx);
      char ch;
      do {
        ch = is.get();
      } while (isspace(ch));
      if (std::is_integral<ValueType>::value && ch == 'L') {
        ch = is.get();
      }
      if (ch == ',') {
        while (true) {
          ch = is.peek();
          if (isspace(ch)) {
            is.get(); continue;
          }
          if (ch == ')' || ch == ']') {
            is.get(); break;
          }
          break;
        }
        if (ch == ')' || ch == ']') break;
      } else if (ch == ')' || ch == ']') {
        break;
      } else {
        is.setstate(std::ios::failbit);
        return is;
      }
    }
    t.assign(tmp.begin(), tmp.end());
    return is;
  }
  template<typename DType = ValueType, typename TStream>
  inline void Save(TStream *strm) const;
  template<typename DType = ValueType, typename TStream>
  inline bool Load(TStream *strm);

 protected:
  // stack cache size
  static const int kStackCache = 4;
  int ndim_{0};
  int num_heap_allocated_{0};
  ValueType data_stack_[kStackCache];
  ValueType* data_heap_{nullptr};
  // internal function to change the dimension
  inline void SetDim(int ndim) {
    CHECK_GE(ndim, -1) << "ndim cannot be less than -1, received " << ndim;
    if (ndim > kStackCache &&
        ndim > num_heap_allocated_) {
      delete [] data_heap_;
      data_heap_ = new ValueType[ndim];
      num_heap_allocated_ = ndim;
    } else if (ndim <= 0 && data_heap_ != nullptr) {
      delete [] data_heap_;
      data_heap_ = nullptr;
      num_heap_allocated_ = 0;
    }
    ndim_ = ndim;
  }
};


inline bool ndim_is_known(const int ndim) {
  CHECK_GE(ndim, -1) << "shape ndim must be >= -1, while received " << ndim;
  return ndim != -1;
}

inline bool dim_size_is_known(const dim_t dim_size) {
  CHECK_GE(dim_size, -1) << "shape dim size must be >= -1, while received " << dim_size;
  return dim_size != -1;
}

class TShape : public Tuple<dim_t> {
 public:
  TShape() {
    this->SetDim(-1);
  }
  inline TShape(const int ndim, const dim_t value) {  // NOLINT(*)
    this->SetDim(ndim);
    if (ndim > 0) {
      std::fill_n(begin(), ndim, value);
    }
  }
  inline TShape(const Tuple<dim_t>& s) { // NOLINT(*)
    if (s.ndim() == -1) {
      this->SetDim(-1);
    } else {
      this->assign(s.begin(), s.end());
    }
  }
  inline TShape(std::initializer_list<dim_t> init) {
    this->assign(init.begin(), init.end());
  }
  inline TShape(Tuple<dim_t>&& s) {  // NOLINT(*)
    this->swap(s);
  }
  template<typename RandomAccessIterator,
           typename std::enable_if<
               std::is_same<typename std::iterator_traits<RandomAccessIterator>::iterator_category,
                            std::random_access_iterator_tag>::value, int>::type = 0>
  inline TShape(RandomAccessIterator begin,
                RandomAccessIterator end) {
    this->assign(begin, end);
  }
  inline TShape& operator=(const Tuple<dim_t>& src) {
    if (src.ndim() == -1) {
      this->SetDim(-1);
    } else {
      this->assign(src.begin(), src.end());
    }
    return *this;
  }
  inline TShape& operator=(Tuple<dim_t>&& src) {  // NOLINT(*)
    TShape(std::move(src)).swap(*this);  // NOLINT(*)
    return *this;
  }
  inline size_t Size() const {
    CHECK(ndim_is_known(this->ndim())) << "Shape is unknown.";
    dim_t size = 1;
    const dim_t* start = begin(), *fin = end();
    for (const dim_t* it = start; it != fin; ++it) {
      CHECK(dim_size_is_known(*it)) << "Shape dim size cannot be a negative value " << *it;
      size *= *it;
    }
    return size;
  }
  inline size_t ProdShape(int dimstart, int dimend) const {
    CHECK(ndim_is_known(this->ndim())) << "Shape is unknown.";
    CHECK_GE(dimstart, 0) << "dimstart must be >= 0, while received " << dimstart;
    CHECK_LE(dimend, this->ndim()) << "dimend must be <= " << this->ndim()
                                   << ", while received " << dimend;
    dim_t num = 1;
    const dim_t *d = this->data();
    for (int i = dimstart; i < dimend; ++i) {
      CHECK(dim_size_is_known(d[i])) << "Shape dim size must be known, while received " << d[i];
      num *= d[i];
    }
    return num;
  }
  inline const dim_t *data() const {
    return begin();
  }
  inline dim_t *data() {
    return begin();
  }
#ifdef MSHADOW_XINLINE
  template<int dim>
  inline TShape(const mshadow::Shape<dim> &s) {// NOLINT(*)
    this->assign(s.shape_, s.shape_ + dim);
  }

  template<int dim>
  inline TShape(mshadow::Shape<dim> &&s) {// NOLINT(*)
    this->assign(s.shape_, s.shape_ + dim);
  }
  template<int dim>
  inline TShape &operator=(const mshadow::Shape<dim> &shape) {
    this->assign(shape.shape_, shape.shape_ + dim);
    return *this;
  }
  template<int dim>
  inline mshadow::Shape<dim> get() const {
    CHECK_EQ(dim, ndim())
        << "dimension do not match target dimension " << dim << " vs " << ndim();
    const dim_t *d = this->data();
    mshadow::Shape<dim> s;
    for (int i = 0; i < dim; ++i) {
      s[i] = d[i];
    }
    return s;
  }
  inline mshadow::Shape<2> FlatTo2D(void) const {
    mshadow::Shape<2> s;
    CHECK(ndim_is_known(ndim())) << "shape must have a valid ndim";
    if (ndim() == 0) return mshadow::Shape2(1, 1);
    const dim_t *d = this->data();
    s.shape_[1] = d[ndim() - 1];
    dim_t ymax = 1;
    for (int i = 1; i < ndim(); ++i) {
      ymax *= d[i - 1];
    }
    s.shape_[0] = ymax;
    return s;
  }
  inline mshadow::Shape<3> FlatTo3D(int axis_begin, int axis_end) const {
    CHECK(axis_end >= axis_begin);
    mshadow::Shape<3> s;
    CHECK(ndim_is_known(ndim())) << "shape must have a valid ndim";
    if (ndim() == 0) return mshadow::Shape3(1, 1, 1);
    const dim_t *d = this->data();
    s.shape_[0] = 1;
    s.shape_[1] = 1;
    s.shape_[2] = 1;

    for (int i = 0; i < axis_begin; ++i) {
      s.shape_[0] *= d[i];
    }
    for (int i = axis_begin; i <= axis_end; ++i) {
      s.shape_[1] *= d[i];
    }
    for (int i = axis_end + 1; i < ndim(); ++i) {
      s.shape_[2] *= d[i];
    }
    return s;
  }
  inline mshadow::Shape<3> FlatTo3D(int axis) const {
    return FlatTo3D(axis, axis);
  }
  inline bool operator==(const TShape &s) const {
    if (ndim() != s.ndim()) return false;
    return std::equal(begin(), end(), s.begin());
  }
  inline bool operator!=(const TShape &s) const {
    return !(*this == s);
  }
  template<int dim>
  inline bool operator==(const mshadow::Shape<dim> &s) const {
    if (ndim_ != dim) return false;
    const dim_t *d = dim <= kStackCache ? data_stack_ : data_heap_;
    for (size_t i = 0; i < dim; ++i) {
      if (d[i] != s.shape_[i]) return false;
    }
    return true;
  }
  template<int dim>
  inline bool operator!=(const mshadow::Shape<dim> &s) const {
    return !(*this == s);
  }
#endif
};

inline bool ndim_is_known(const TShape& x) {
  return ndim_is_known(x.ndim());
}

inline bool dim_size_is_known(const TShape& x, const int idx) {
  CHECK(idx >= 0 && idx < x.ndim())
      << "idx = " << idx << " exceeds shape dimension range [0, " << x.ndim() << ")";
  return dim_size_is_known(x[idx]);
}

inline bool shape_is_known(const TShape& x) {
  if (!ndim_is_known(x)) return false;
  for (int i = 0; i < x.ndim(); ++i) {
    if (!dim_size_is_known(x, i)) return false;
  }
  return true;
}

template<typename SrcIter, typename DstIter>
inline DstIter ShapeTypeCast(const SrcIter begin,
                             const SrcIter end,
                             DstIter dst_begin) {
  typedef typename std::iterator_traits<SrcIter>::value_type SrcDType;
  typedef typename std::iterator_traits<DstIter>::value_type DstDType;
  auto cast = [](const SrcDType& dim) { return static_cast<DstDType>(dim); };
  return std::transform(begin, end, dst_begin, cast);
}

template<typename SrcIter>
inline TShape ShapeTypeCast(const SrcIter begin, const SrcIter end) {
  size_t ndim = std::distance(begin, end);
  TShape res(ndim, -1);
  ShapeTypeCast(begin, end, res.begin());
  return res;
}

template<typename ValueType>
template<typename DType, typename TStream>
inline void Tuple<ValueType>::Save(TStream *strm) const {
  strm->Write(&ndim_, sizeof(ndim_));
  if (typeid(DType) == typeid(ValueType)) {
    strm->Write(begin(), sizeof(ValueType) * ndim_);
  } else {
    std::vector<DType> buffer(ndim_);
    ShapeTypeCast(begin(), end(), buffer.data());
    strm->Write(buffer.data(), sizeof(DType) * ndim_);
  }
}

template<typename ValueType>
template<typename DType, typename TStream>
inline bool Tuple<ValueType>::Load(TStream *strm) {
  if (strm->Read(&ndim_, sizeof(ndim_)) != sizeof(ndim_)) return false;
  this->SetDim(ndim_);
  size_t nread = sizeof(DType) * ndim_;
  if (typeid(DType) == typeid(ValueType)) {
    if (strm->Read(begin(), nread) != nread) return false;
  } else {
    std::vector<DType> buffer(ndim_);
    if (strm->Read(buffer.data(), nread) != nread) return false;
    ShapeTypeCast(buffer.begin(), buffer.end(), begin());
  }
  return true;
}

}  // namespace mxnet

namespace std {
template<typename T>
struct hash<mxnet::Tuple<T> > {
  size_t operator()(const mxnet::Tuple<T>& val) const {
    std::hash<int> hash_int;
    size_t res = hash_int(val.ndim());
    for (int i = 0; i < val.ndim(); ++i) {
      res = dmlc::HashCombine(res, val[i]);
    }
    return res;
  }
};

template<>
struct hash<mxnet::TShape> {
  size_t operator()(const mxnet::TShape& val) const {
    std::hash<int> hash_int;
    size_t res = hash_int(val.ndim());
    for (int i = 0; i < val.ndim(); ++i) {
      res = dmlc::HashCombine(res, val[i]);
    }
    return res;
  }
};
}  // namespace std

namespace dmlc {
DMLC_DECLARE_TYPE_NAME(optional<mxnet::TShape>, "Shape or None");
DMLC_DECLARE_TYPE_NAME(optional<mxnet::Tuple<int>>, "Shape or None");
// avoid low version of MSVC
#if !defined(_MSC_VER)
template<typename T>
struct type_name_helper<mxnet::Tuple<T> > {
  static inline std::string value() {
    return "tuple of <" + type_name<T>() + ">";
  }
};
#endif
}  // namespace dmlc

namespace mxnet {
using ShapeVector = std::vector<mxnet::TShape>;

using FInferShape = nnvm::FInferNodeEntryAttr<mxnet::TShape>;

}  // namespace mxnet

#endif  // MXNET_TUPLE_H_