ndarray.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_NDARRAY_H_
#define MXNET_NDARRAY_H_

#include <dmlc/base.h>
#include <dmlc/logging.h>
#include <dmlc/io.h>
#include <dmlc/type_traits.h>
#include <dmlc/registry.h>
#include <nnvm/node.h>
#include <vector>
#include <map>
#include <string>
#include <algorithm>
#include <memory>
#include <algorithm>
#if MXNET_USE_MKLDNN == 1
#include <mkldnn.hpp>
#endif
#include "./base.h"
#include "./storage.h"
#include "./engine.h"
// check c++11
#if DMLC_USE_CXX11 == 0
#error "cxx11 was required for ndarray module"
#endif

namespace mxnet {
// enum for storage types
namespace csr {
enum CSRAuxType {kIndPtr, kIdx};
}

namespace rowsparse {
enum RowSparseAuxType {kIdx};
}

enum NDArrayStorageType {
  kUndefinedStorage = -1,  // undefined storage
  kDefaultStorage,         // dense
  kRowSparseStorage,       // row sparse
  kCSRStorage,             // csr
};

enum NDArrayFormatErr {
  kNormalErr,     // normal
  kCSRShapeErr,   // shape mismatch for csr
  kCSRIndPtrErr,  // indptr error for csr
  kCSRIdxErr,     // idx error for csr
  kRSPShapeErr,   // shape mismatch for row sparse
  kRSPIdxErr,     // indices error for row sparse
};

class MKLDNNMemory;

class NDArray {
 public:
  NDArray() {
  }
  NDArray(const TShape &shape, Context ctx,
          bool delay_alloc = false, int dtype = mshadow::default_type_flag)
      : ptr_(std::make_shared<Chunk>(shape, ctx, delay_alloc, dtype)),
        shape_(shape), dtype_(dtype), storage_type_(kDefaultStorage),
        entry_({nullptr, 0, 0}) {
  }
  NDArray(const NDArrayStorageType stype, const TShape &shape, Context ctx,
          bool delay_alloc = true, int dtype = mshadow::default_type_flag,
          std::vector<int> aux_types = {}, std::vector<TShape> aux_shapes = {},
          TShape storage_shape = TShape(mshadow::Shape1(0)));
  explicit NDArray(Context ctx, int dtype = mshadow::default_type_flag) {
    ptr_ = std::make_shared<Chunk>(TShape(mshadow::Shape1(0)), ctx, true, dtype);
    dtype_ = dtype;
    storage_type_ = kDefaultStorage;
    entry_ = {nullptr, 0, 0};
  }
  NDArray(const TBlob &data, int dev_id)
      : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_),
        dtype_(data.type_flag_), storage_type_(kDefaultStorage),
        entry_({nullptr, 0, 0}) {
  }

  NDArray(const TBlob &data, int dev_id, const std::function<void()>& deleter)
      : ptr_(new Chunk(data, dev_id),
        [deleter](Chunk *p) {
          deleter();    // call custom deleter
          delete p;     // delete Chunk object
        }),
        shape_(data.shape_),
        dtype_(data.type_flag_), storage_type_(kDefaultStorage),
        entry_({nullptr, 0, 0}) {
  }

  NDArray(int shared_pid, int shared_id, const TShape& shape, int dtype)
      : ptr_(std::make_shared<Chunk>(shared_pid, shared_id, shape, dtype)), shape_(shape),
        dtype_(dtype), storage_type_(kDefaultStorage), entry_({nullptr, 0, 0}) {
  }

  NDArray(const NDArrayStorageType stype, const TShape &shape,
          const TBlob &data, const std::vector<TBlob> &aux_data, int dev_id)
      : ptr_(std::make_shared<Chunk>(stype, data, aux_data, dev_id)), shape_(shape),
        dtype_(data.type_flag_), storage_type_(stype), entry_({nullptr, 0, 0}) {
  }
  void Init(const TShape &shape) {
    ptr_->Init(shape, this->dtype_);
    this->shape_ = shape;
  }
  void SetShapeFromChunk() {
    shape_ = ptr_->storage_shape;
  }
  /*
   * This indicates whether an array is a view of another array (created by
   * reshape or slice). If an array is a view and the the data is stored in
   * MKLDNN format, we need to convert the data to the default format when
   * data in the view is accessed.
   */
  inline bool IsView() const {
    // View only works on the default storage
    if (storage_type() != kDefaultStorage)
      return false;
    // If the array reuses memory, its shape may be different from the storage
    // shape. However, we shouldn't consider it as a view.
    if (reuse_)
      return false;
    return byte_offset_ > 0 || shape() != ptr_->storage_shape;
  }

  /* \brief Check whether the two arrays are the same array */
  inline bool IsSame(const NDArray& other) const {
    return ptr_ == other.ptr_ &&
        shape_ == other.shape_ &&
        byte_offset_ == other.byte_offset_ &&
        dtype_ == other.dtype_;
  }

  inline const TShape& shape() const {
    return shape_;
  }
  inline const TShape &storage_shape() const {
    CHECK(ptr_ != nullptr);
    CHECK_NE(storage_type(), kDefaultStorage)
             << "storage_shape() is not intended for kDefaultStorage.";
    return ptr_->storage_shape;
  }

  inline const TShape& aux_shape(size_t index) const {
    CHECK_NE(storage_type(), kDefaultStorage)
             << "aux_shape() is not intended for kDefaultStorage.";
    return ptr_->aux_shapes[index];
  }

  /* \return the shapes of all aux data */
  const std::vector<TShape>& aux_shapes() const {
    CHECK_NE(storage_type(), kDefaultStorage)
             << "aux_shapes() is not intended for kDefaultStorage.";
    return ptr_->aux_shapes;
  }

  const std::vector<int>& aux_types() const {
    CHECK_NE(storage_type(), kDefaultStorage)
      << "aux_types() is not intended for kDefaultStorage.";
    return ptr_->aux_types;
  }

  inline void set_aux_shape(size_t index, const TShape& shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
      << "set_aux_shape() is not intended for kDefaultStorage.";
    ptr_->set_aux_shape(index, shape);
  }

  inline const TBlob& data() const {
    if (storage_type() == kDefaultStorage) CheckAndAlloc();
    SetTBlob();
    return tblob_;
  }
  NDArray grad() const;

  inline TBlob aux_data(size_t i) const {
    auto stype = storage_type();
    TBlob res;
    auto shape = aux_shape(i);
    auto type = aux_type(i);
    MSHADOW_TYPE_SWITCH(type, DType, {
      auto dptr = static_cast<DType*>(ptr_->aux_handles[i].dptr);
      CHECK(stype == kRowSparseStorage || stype == kCSRStorage)
            << "Unexpected storage type: " << stype;
      res = TBlob(dptr, shape, ptr_->aux_handles[i].ctx.dev_mask(), type);
    });
    return res;
  }
  inline Context ctx() const {
    CHECK(!is_none());
    return ptr_->shandle.ctx;
  }
  inline int dtype() const {
    return dtype_;
  }
  inline int aux_type(size_t i) const {
    CHECK(!is_none());
    return ptr_->aux_types[i];
  }

  inline NDArrayStorageType storage_type() const {
    return storage_type_;
  }
  inline bool is_none() const {
    return ptr_.get() == nullptr;
  }
  bool fresh_out_grad() const;
  void set_fresh_out_grad(bool state) const;
  inline bool storage_initialized() const {
    if (is_none()) return false;
    auto stype = storage_type();
    CHECK_NE(stype, kDefaultStorage)
             << "storage_initialized() is not intended for kDefaultStorage.";
    if (stype == kRowSparseStorage) {
      CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
               << "inconsistent storage shape " << storage_shape()
               << " vs. aux shape " << aux_shape(rowsparse::kIdx);
      return aux_shape(rowsparse::kIdx).Size() != 0;
    } else if (stype == kCSRStorage) {
      CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
               << "inconsistent storage shape " << storage_shape()
               << " vs. aux shape " << aux_shape(csr::kIdx);
      return aux_shape(csr::kIdx).Size() != 0;
    } else {
      LOG(FATAL) << "Unknown storage type";
    }
    return true;
  }
  inline Storage::Handle storage_handle() const {
    CHECK(!is_none());
    CHECK_EQ(storage_type(), kDefaultStorage);
    CheckAndAlloc();
    return ptr_->shandle;
  }
  inline void WaitToRead() const {
    if (is_none()) return;
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline void WaitToWrite() const {
    if (is_none()) return;
    Engine::Get()->PushAsync(
      [](RunContext, Engine::CallbackOnComplete on_complete) {
        on_complete();
      }, Context{}, {}, {ptr_->var});
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline Engine::VarHandle var() const {
    return ptr_->var;
  }
  inline size_t byte_offset() const {
    return byte_offset_;
  }
  inline size_t version() const {
    return var()->version();
  }
  void Save(dmlc::Stream *strm) const;
  bool LegacyLoad(dmlc::Stream *strm, const uint32_t magic);
  bool Load(dmlc::Stream *strm);
  NDArray &operator=(real_t scalar);
  NDArray &operator+=(const NDArray &src);
  NDArray &operator+=(const real_t &src);
  NDArray &operator-=(const NDArray &src);
  NDArray &operator-=(const real_t &src);
  NDArray &operator*=(const NDArray &src);
  NDArray &operator*=(const real_t &src);
  NDArray &operator/=(const NDArray &src);
  NDArray &operator/=(const real_t &src);
  NDArray Copy(Context ctx) const;
  void SyncCopyFromCPU(const void *data, size_t size) const;

  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);

  void SyncCopyToCPU(void *data, size_t size) const;
  void SyncCheckFormat(const bool full_check) const;
  NDArray Slice(index_t begin, index_t end) const;
  NDArray SliceWithRecord(index_t begin, index_t end);
  NDArray At(index_t idx) const;
  NDArray AtWithRecord(index_t idx);
  NDArray aux_ndarray(size_t i) const;

  NDArray data_ndarray() const;

  inline NDArray AsArray(const TShape &shape, int dtype) const {
    CHECK_EQ(storage_type(), kDefaultStorage)
             << "AsArray is intended only for kDefaultStorage.";
    CHECK_GE(ptr_->shandle.size,
             shape.Size() * mshadow::mshadow_sizeof(dtype))
        << "NDArray.AsArray: target memory size is bigger";
    // We can't reuse memory in a view.
    CHECK(!IsView());
    NDArray ret = *this;
    ret.shape_ = shape;
    ret.dtype_ = dtype;
    ret.reuse_ = true;
    return ret;
  }

  DLManagedTensor* ToDLPack() const;

  static NDArray FromDLPack(const DLManagedTensor* tensor);

  inline void SparseUpdateChunk(const NDArray &arr) const {
    CHECK(shape_ == arr.shape_) << "ndarray shape is different from the target";
    CHECK(dtype_ == arr.dtype_) << "ndarray dtype is different from the target";
    auto stype = arr.storage_type();
    CHECK(stype == kCSRStorage || stype == kRowSparseStorage)
        << "Only to be used with CSR and RSP storage types";
    // swap shandles between src and dst
    Storage::Handle shandle_dst = arr.ptr_->shandle;
    arr.ptr_->shandle = ptr_->shandle;
    ptr_->shandle = shandle_dst;

    ptr_->storage_shape = arr.ptr_->storage_shape;
    ptr_->storage_type = arr.ptr_->storage_type;
    ptr_->ctx = arr.ptr_->ctx;

    // swap aux_handles between src and dst
    size_t aux_idx = 0;
    CHECK(ptr_->aux_handles.size() == arr.ptr_->aux_handles.size())
        << "ndarray number of aux_handles is different from target";
    for (auto &aux_handle : arr.ptr_->aux_handles) {
      Storage::Handle aux_dst = ptr_->aux_handles[aux_idx];
      ptr_->aux_handles[aux_idx] = aux_handle;
      aux_handle = aux_dst;
      aux_idx++;
    }
    ptr_->aux_types = arr.ptr_->aux_types;
    ptr_->aux_shapes = arr.ptr_->aux_shapes;
  }

  NDArray Reshape(const TShape &shape) const;
  NDArray ReshapeWithRecord(const TShape &shape);
  NDArray Detach() const {
    NDArray ret(*this);
    ret.entry_ = nnvm::NodeEntry{nullptr, 0, 0};
    return ret;
  }

  nnvm::Symbol get_autograd_symbol() const;
  inline void CheckAndAlloc() const {
    CHECK_EQ(storage_type(), kDefaultStorage);
    ptr_->CheckAndAlloc();
  }

  void ReshapeAndAlloc(const TShape& shape) {
    CHECK_EQ(storage_type(), kDefaultStorage);
    CHECK(!is_none());
    shape_ = shape;
    ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
  }

  /* !
   * \brief Alloc memory for non-default storage
   * aux_shape is only known at run time
   */
  inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) const {
    CHECK_NE(storage_type(), kDefaultStorage)
             << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
    ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
  }
  inline void CheckAndAllocData(const TShape &storage_shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
             << "CheckAndAllocData is not intended for kDefaultStorage";
    ptr_->CheckAndAllocData(storage_shape, dtype_);
  }
  inline void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
             << "CheckAndAllocAuxData is not intended for kDefaultStorage";
    ptr_->CheckAndAllocAuxData(i, aux_shape);
  }

#if MXNET_USE_MKLDNN == 1
  /*
   * Create NDArray from mkldnn memory.
   * mkldnn_mem The mkldnn memory to be managed.
   */
  explicit NDArray(const std::shared_ptr<mkldnn::memory> &mkldnn_mem);
  /*
   * Create NDArray from mkldnn memory descriptor.
   * mem_pd The mkldnn memory descriptor to be created.
   */
  explicit NDArray(mkldnn::memory::primitive_desc mem_pd);
  /*
   * Test if the data is stored in one of special MKLDNN format.
   */
  bool IsMKLDNNData() const {
    return ptr_->IsMKLDNN();
  }
  /*
   * Test if the data is stored in one of default MXNet formats.
   */
  bool IsDefaultData() const {
    return ptr_->IsDefault();
  }
  /*
   * All functions below return a raw pointer to mkldnn memory. Actually there
   * is a shared pointer that hold the memory either in NDArray or in MKLDNN
   * stream. As long as we call these functions inside an operator, the return
   * memory is always valid.
   */

  /*
   * This function returns mkldnn::memory with the default primitive_desc.
   */
  const mkldnn::memory *GetMKLDNNData() const;
  /*
   * This function returns mkldnn::memory with the given primitive_desc
   * as long as the array size meets the required size in the given primitive_desc.
   */
  const mkldnn::memory *GetMKLDNNData(
      const mkldnn::memory::primitive_desc &desc) const;
  /*
   * This function returns mkldnn::memory with the given primitive_desc.
   * The returned mkldnn::memory will have the same physical layout as
   * the given primitive_desc.
   */
  const mkldnn::memory *GetMKLDNNDataReorder(
      const mkldnn::memory::primitive_desc &desc) const;

  /*
   * This function copies data from mkldnn memory.
   */
  void CopyFrom(const mkldnn::memory &mem);
  /*
   * This function allocates memory for array and creates mkldnn memory
   * with the specified format.
   */
  mkldnn::memory *CreateMKLDNNData(
      const mkldnn::memory::primitive_desc &desc);

  /*
   * These are the async version of the methods above.
   * It changes the layout of this NDArray, but it happens after all accesses to
   * the array are complete.
   */
  void Reorder2DefaultAsync();
  void MKLDNNDataReorderAsync(const mkldnn::memory::primitive_desc &desc);

  /*
   * This creates a new NDArray with the reordered data.
   * It doesn't affect the data of the original NDArray.
   */
  NDArray Reorder2Default() const;

  void InvalidateMKLDNNData();

  /*
   * This function is used inside operators to reshape an array.
   * It doesn't change the layout of the original array and allocate memory from
   * the temporary buffer. The returned array is only valid inside the current
   * invocation of this operator.
   * This is different from Reshape. Reshape will cause data in the array to be
   * converted to the default layout and allocate memory from malloc directly,
   * which can be expensive.
   * It's used by FullyConnected right now.
   */
  NDArray MKLDNNDataReshape(const TShape &shape) const;

  void UpdateMKLDNNMemDesc(mkldnn::memory::format format);
#endif

  static void Save(dmlc::Stream* fo,
                   const std::vector<NDArray>& data,
                   const std::vector<std::string>& names);
  static void Load(dmlc::Stream* fi,
                   std::vector<NDArray>* data,
                   std::vector<std::string>* keys);

 private:
  friend class Imperative;
  // shandle is used to store the actual values in the NDArray
  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
  struct Chunk {
    Storage::Handle shandle;
    std::vector<Storage::Handle> aux_handles;

#if MXNET_USE_MKLDNN == 1

    std::shared_ptr<MKLDNNMemory> mkl_mem_;
#endif

    Engine::VarHandle var;
    bool static_data;
    bool delay_alloc;
    // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
    // is constructed.
    NDArrayStorageType storage_type = kDefaultStorage;
    std::vector<int> aux_types;
    // context of data
    Context ctx;
    // The shape of the chunk data.
    // This might not be the same shape as the NDArray, since the storage may be sparse.
    // The default value for storage_shape is {0} when an empty non-default NDArray is created.
    TShape storage_shape;
    // The shape of aux data. The default value for the shape depends on the type of storage.
    // If aux_shapes[i].Size() is zero, aux data i is empty.
    std::vector<TShape> aux_shapes;

    Chunk() : static_data(true), delay_alloc(false) {}

    Chunk(TShape shape, Context ctx_, bool delay_alloc_, int dtype)
        : static_data(false), delay_alloc(true), ctx(ctx_) {
      auto size = shape.Size();
      storage_shape = shape;
      var = Engine::Get()->NewVariable();
      shandle.size = size * mshadow::mshadow_sizeof(dtype);
      shandle.ctx = ctx_;
      if (!delay_alloc_) this->CheckAndAlloc();
    }

    Chunk(const TBlob &data, int dev_id)
        : static_data(true), delay_alloc(false) {
      CHECK(storage_type == kDefaultStorage);
      var = Engine::Get()->NewVariable();
      if (data.dev_mask() == cpu::kDevMask) {
        ctx = Context::CPU();
      } else {
        CHECK_EQ(data.dev_mask(), gpu::kDevMask);
        ctx = Context::GPU(dev_id);
      }
      // init shandle
      shandle.ctx = ctx;
      shandle.dptr = data.dptr_;
      shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
      storage_shape = data.shape_;
    }

    Chunk(int shared_pid, int shared_id, const TShape& shape, int dtype)
        : static_data(false), delay_alloc(false) {
      var = Engine::Get()->NewVariable();
      ctx = Context::CPUShared(0);
      shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);
      shandle.ctx = ctx;
      shandle.shared_pid = shared_pid;
      shandle.shared_id = shared_id;
      Storage::Get()->Alloc(&shandle);
      storage_shape = shape;
    }
    // Constructor for a non-default storage chunk
    Chunk(NDArrayStorageType storage_type_, const TShape &storage_shape_, Context ctx_,
          bool delay_alloc_, int dtype, const std::vector<int> &aux_types_,
          const std::vector<TShape> &aux_shapes_)
        : static_data(false), delay_alloc(delay_alloc_), storage_type(storage_type_),
          aux_types(aux_types_), ctx(ctx_), storage_shape(storage_shape_),
          aux_shapes(aux_shapes_) {
      shandle.ctx = ctx;
      var = Engine::Get()->NewVariable();
      // aux_handles always reflect the correct number of aux data
      for (size_t i = 0; i < aux_shapes.size(); i++) {
        CheckAndAllocAuxData(i, aux_shapes[i]);
        // this line is needed in case when aux_shapes[i].Size() = 0
        // aux_handles[i] will not be updated and take only default value.
        aux_handles[i].ctx = ctx;
      }
      if (!delay_alloc) {
        CheckAndAllocData(storage_shape, dtype);
      }
    }

    Chunk(const NDArrayStorageType storage_type_, const TBlob &data,
          const std::vector<TBlob> &aux_data, int dev_id)
        : static_data(true), delay_alloc(false), storage_type(storage_type_) {
      using namespace mshadow;
      CHECK_NE(storage_type, kDefaultStorage);
      // init var
      var = Engine::Get()->NewVariable();
      // init ctx
      if (data.dev_mask() == cpu::kDevMask) {
        ctx = Context::CPU();
      } else {
        CHECK_EQ(data.dev_mask(), gpu::kDevMask);
        ctx = Context::GPU(dev_id);
      }
      // init shandle
      shandle.ctx = ctx;
      shandle.dptr = data.dptr_;
      shandle.size = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
      storage_shape = data.shape_;
      // init aux handles
      for (const auto &aux : aux_data) {
        Storage::Handle aux_handle;
        aux_handle.ctx = ctx;
        aux_handle.dptr = aux.dptr_;
        aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
        aux_handles.push_back(aux_handle);
        aux_types.emplace_back(aux.type_flag_);
        aux_shapes.emplace_back(aux.shape_);
      }
    }

    inline void set_aux_shape(const size_t i, const TShape& shape) {
      aux_shapes[i] = shape;
      if (storage_shape.ndim() > 0) {
        if (storage_type == kRowSparseStorage && i == rowsparse::kIdx) {
          storage_shape[0] = shape[0];
        } else if (storage_type == kCSRStorage && i == csr::kIdx) {
          storage_shape[0] = shape[0];
        }
      }
    }

    inline void CheckAndAlloc(void) {
      if (delay_alloc) {
        shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
#if MXNET_USE_MKLDNN == 1
        mkl_mem_ = nullptr;
#endif
        delay_alloc = false;
      }
    }

    // size is the number of bytes
    void CheckAndAlloc(uint64_t dbytes) {
      CHECK_EQ(kDefaultStorage, storage_type)
          << "CheckAndAlloc(dbytes) is only intended for kDefaultStorage";
      dbytes = std::max(dbytes, static_cast<uint64_t>(shandle.size));
      if (delay_alloc) {
        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
#if MXNET_USE_MKLDNN == 1
        mkl_mem_ = nullptr;
#endif
        delay_alloc = false;
      } else if (shandle.size < dbytes) {
        // free storage
        Storage::Get()->Free(shandle);
        // init storage
        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
#if MXNET_USE_MKLDNN == 1
        mkl_mem_ = nullptr;
#endif
      }
    }
    void Init(const TShape &shape, int dtype) {
      auto size = shape.Size();
      storage_shape = shape;
      shandle.size = size * mshadow::mshadow_sizeof(dtype);
      this->CheckAndAlloc();
    }
    inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
                              int dtype) {
      // calculate size, perform allocation
      if (kRowSparseStorage == storage_type) {
        // For row sparse, aux_shape indicates the number of rows to allocate
        auto aux_shape = aux_shapes[rowsparse::kIdx];
        CheckAndAllocAuxData(rowsparse::kIdx, aux_shape);
        TShape storage_shape(shape);
        storage_shape[0] = aux_shape[0];
        CheckAndAllocData(storage_shape, dtype);
      } else if (kCSRStorage == storage_type) {
        CheckAndAllocAuxData(csr::kIndPtr, aux_shapes[csr::kIndPtr]);
        CheckAndAllocAuxData(csr::kIdx, aux_shapes[csr::kIdx]);
        CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
      } else {
        LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
      }
    }
    // create storage handle for data based on shape and dtype, assuming ctx is set
    // storage shape is also updated
    // if data is already allocated, try reuse the storage. Otherwise, free the current one
    // and allocate new storage
    void CheckAndAllocData(const TShape &shape, int dtype);

#if MXNET_USE_MKLDNN == 1
    // Have MKL memory reference to the data in the default storage
    // or create memory for MKLDNN.
    void SetMKLMem(const TShape &shape, int dtype);
    // If the data is stored in MKLDNN layout, we reorder data in mkl_mem_ and
    // save the result in shandle.
    void Reorder2Default();
    // Reroder data to a specified layout.
    void MKLDNNDataReorder(const mkldnn::memory::primitive_desc &desc);
    bool IsMKLDNN() const;
    bool IsDefault() const;
#endif

    // create storage handle for aux data based on shape
    // this function assumes ctx, aux shapes and aux types are set
    // aux shape is also updated
    // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
    // and allocate new storage
    inline void CheckAndAllocAuxData(size_t i, const TShape &shape) {
      CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
      CHECK_NE(storage_type, kUndefinedStorage)
        << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
      CHECK_NE(storage_type, kDefaultStorage)
        << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
      if (aux_handles.size() <= i) {
        aux_handles.resize(i + 1);
      }
      size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
      if (aux_handles[i].size < aux_bytes) {
        // free storage
        Storage::Get()->Free(aux_handles[i]);
        // init aux storage
        aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
      }
      // init shape
      set_aux_shape(i, shape);
    }
    ~Chunk();
  };  // struct Chunk

  void SetTBlob() const;

  std::shared_ptr<Chunk> ptr_{nullptr};
  TShape shape_;
  size_t byte_offset_ = 0;
  int dtype_ = -1;
  bool reuse_ = false;
  NDArrayStorageType storage_type_ = kUndefinedStorage;
  nnvm::NodeEntry entry_;
  mutable TBlob tblob_;
};  // class NDArray

size_t num_aux_data(NDArrayStorageType stype);

void CopyFromTo(const NDArray &from, const NDArray *to, int priority = 0);

void CopyFromTo(const NDArray &from, const NDArray& to, int priority = 0, bool is_opr = false);

void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);

NDArray operator+(const NDArray &lhs, const NDArray &rhs);
NDArray operator+(const NDArray &lhs, const real_t &rhs);
NDArray operator-(const NDArray &lhs, const NDArray &rhs);
NDArray operator-(const NDArray &lhs, const real_t &rhs);
NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
NDArray operator*(const NDArray &lhs, const real_t &rhs);
NDArray operator/(const NDArray &lhs, const NDArray &rhs);
NDArray operator/(const NDArray &lhs, const real_t &rhs);

void RandomSeed(uint32_t seed);
void RandomSeed(Context ctx, uint32_t seed);
void SampleUniform(real_t begin, real_t end, NDArray *out);
void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
void SampleGamma(real_t alpha, real_t beta, NDArray *out);
void SampleExponential(real_t lambda, NDArray *out);
void SamplePoisson(real_t lambda, NDArray *out);
void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);


//--------------------------------------------------------------
// The following part are API Registration of NDArray functions.
//--------------------------------------------------------------

typedef std::function<void (NDArray **used_vars,
                            real_t *scalars,
                            NDArray **mutate_vars,
                            int num_params,
                            char **param_keys,
                            char **param_vals)> NDArrayAPIFunction;
enum NDArrayFunctionTypeMask {
  kNDArrayArgBeforeScalar = 1,
  kScalarArgBeforeNDArray = 1 << 1,
  kAcceptEmptyMutateTarget = 1 << 2
};
struct NDArrayFunctionReg
    : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
                                        NDArrayAPIFunction> {
  unsigned num_use_vars;
  unsigned num_mutate_vars;
  unsigned num_scalars;
  int type_mask;
  NDArrayFunctionReg()
      : num_use_vars(0),
        num_mutate_vars(0),
        num_scalars(0),
        type_mask(0) {}
  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
                                                            NDArray *out)) {
    body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                        int num_params, char **param_keys, char **param_vals) {
      (*fsetvalue)(s[0], mutate_vars[0]);
    };
    num_mutate_vars = 1; num_scalars = 1;
    this->add_argument("src", "real_t", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
                                                          const NDArray &mhs,
                                                          const NDArray &rhs,
                                                                NDArray *out)) {
    body = [fternary](NDArray **used_vars,
      real_t *s, NDArray **mutate_vars,
      int num_params, char **param_keys, char **param_vals) {
      (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
    };
    num_use_vars = 3; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("mhs", "NDArray", "Middle operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
                                                          const NDArray &rhs,
                                                          NDArray *out)) {
    body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                      int num_params, char **param_keys, char **param_vals) {
      (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
    };
    num_use_vars = 2; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
                                                          const real_t &rhs,
                                                          NDArray *out)) {
    body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                      int num_params, char **param_keys, char **param_vals) {
      (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
    };
    num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "real_t", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
                                                         NDArray *out)) {
    body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                     int num_params, char **param_keys, char **param_vals) {
      (*funary)(*used_vars[0], mutate_vars[0]);
    };
    num_use_vars = 1; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("src", "NDArray", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(
    void (*fgeneric)(NDArray **used_vars,
                     real_t *s,
                     NDArray **mutate_vars,
                     const std::map<std::string, std::string>& param)) {
    body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                       int num_params, char **param_keys, char **param_vals) {
      std::map<std::string, std::string> param;
      for (int i = 0; i < num_params; ++i) {
        param[param_keys[i]] = param_vals[i];
      }
      fgeneric(used_vars, s, mutate_vars, param);
    };
    return *this;
  }
  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
    num_use_vars = n; return *this;
  }
  inline NDArrayFunctionReg &set_num_mutate_vars(unsigned n) {
    num_mutate_vars = n; return *this;
  }
  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
    num_scalars = n; return *this;
  }
  inline NDArrayFunctionReg &set_type_mask(int tmask) {
    type_mask = tmask; return *this;
  }
};  // NDArrayFunctionReg

#define MXNET_REGISTER_NDARRAY_FUN(name)                                 \
  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)

}  // namespace mxnet

namespace dmlc {
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
}  // namespace dmlc
#endif  // MXNET_NDARRAY_H_