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 <memory>
#include "./base.h"
#include "./storage.h"
#include "./engine.h"
#if MKL_EXPERIMENTAL == 1
#include <mkl_memory.h>
#endif
// 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
};


class NDArray {
 public:
  NDArray() {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = MKLMemHolder::create();
#endif
  }
  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}) {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }
  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)))
      : shape_(shape), dtype_(dtype), storage_type_(stype),
        entry_({nullptr, 0, 0}) {
      // Assign default aux types if not given
      if (aux_types.size() == 0) {
        if (stype == kRowSparseStorage) {
          aux_types = {mshadow::kInt64};
        } else if (stype == kCSRStorage) {
          aux_types = {mshadow::kInt64, mshadow::kInt64};
        } else {
          LOG(FATAL) << "Unknown storage type " << stype;
        }
      }
      // Assign default shapes if not given
      // unknown shapes are intialized as {0} such that Size() would return 0
      if (aux_shapes.size() == 0) {
        if (stype == kRowSparseStorage) {
          aux_shapes = {TShape(mshadow::Shape1(0))};
        } else if (stype == kCSRStorage) {
          // aux shapes for indptr and indices
          aux_shapes = {TShape(mshadow::Shape1(0)), TShape(mshadow::Shape1(0))};
        } else {
          LOG(FATAL) << "Unknown storage type " << stype;
        }
      }
      if (storage_shape.Size() == 0) {
        if (stype == kRowSparseStorage) {
          storage_shape = shape;
          storage_shape[0] = aux_shapes[rowsparse::kIdx][0];
        } else if (stype == kCSRStorage) {
          storage_shape = aux_shapes[csr::kIdx];
        } else {
          LOG(FATAL) << "Unknown storage type " << stype;
        }
      }
      ptr_ = std::make_shared<Chunk>(stype, storage_shape, ctx, delay_alloc,
                                     dtype, aux_types, aux_shapes);
#if MKL_EXPERIMENTAL == 1
      Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }
  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}) {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }

  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}) {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }


  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 {
    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);
    });
#if MKL_EXPERIMENTAL == 1
    res.Mkl_mem_ = Mkl_mem_;
#endif
    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;
  // returns true if a sparse ndarray's aux_data and storage are initialized
  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(0).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(0).Size() != 0;
    } else {
      LOG(FATAL) << "Unknown storage type";
    }
    return true;
  }
  inline void WaitToRead() const {
    if (is_none()) return;
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline void WaitToWrite() const {
    if (is_none()) return;
    Engine::Get()->PushSync([](RunContext) {}, Context{}, {}, {ptr_->var});
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline Engine::VarHandle var() const {
    return ptr_->var;
  }
  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;
  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";
#if MKL_EXPERIMENTAL == 1
    if (Mkl_mem_ != nullptr) {
      // convert prv to cpu
      Mkl_mem_->check_and_prv_to_cpu(ptr_->shandle.dptr);
    }
#endif
    NDArray ret = *this;
    ret.shape_ = shape;
    ret.dtype_ = dtype;
    return ret;
  }
  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);
  }
  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;
    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_;
    }
    // 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);
        delay_alloc = false;
      }
    }

    // size is the number of bytes
    void CheckAndAlloc(uint64_t dbytes) {
      CHECK_EQ(kDefaultStorage, storage_type)
              << "CheckAndAlloc(dbytes) is not intended for kDefaultStorage";
      if (delay_alloc) {
        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
        delay_alloc = false;
      } else if (shandle.size < dbytes) {
        // free storage if necessary and alloc again
        if (shandle.size > 0) Storage::Get()->Free(shandle);
        // init storage
        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
      }
    }

    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
    inline void CheckAndAllocData(const TShape &shape, int dtype) {
      CHECK_NE(aux_shapes.size(), 0) << "data is expected to be allocated after aux_data";
      auto dbytes = shape.Size() * mshadow::mshadow_sizeof(dtype);
      if (shandle.size < dbytes) {
        // free storage if necessary and alloc again
        if (shandle.size > 0) Storage::Get()->Free(shandle);
        // init storage
        shandle = Storage::Get()->Alloc(dbytes, ctx);
      }
      // init shape
      storage_shape = shape;
      // delay_alloc is only set when data storage handle is present
      delay_alloc = false;
    }
    // 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 if necessary and alloc again
        if (aux_handles[i].size > 0) 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() {
      bool skip_free = static_data || delay_alloc;
      Storage::Handle h = this->shandle;
      std::vector<Storage::Handle> aux_h = this->aux_handles;
      Engine::Get()->DeleteVariable([h, aux_h, skip_free](RunContext s) {
        if (skip_free == false) {
          Storage::Get()->Free(h);
          for (size_t i = 0; i < aux_h.size(); i++) {
            if (aux_h[i].size > 0) Storage::Get()->Free(aux_h[i]);
          }
        }
      }, shandle.ctx, var);
    }
  };  // struct Chunk

  void SetTBlob() const {
    CHECK(ptr_ != nullptr);
    TShape shape = shape_;
    char *dptr = static_cast<char*>(ptr_->shandle.dptr);
    auto stype = storage_type();
    if (stype == kDefaultStorage) {
      dptr += byte_offset_;
    } else if (stype == kCSRStorage || stype == kRowSparseStorage) {
      shape = storage_shape();
    } else {
      LOG(FATAL) << "unknown storage type " << stype;
    }
    tblob_.dptr_ = dptr;
    tblob_.shape_ = shape;
    tblob_.type_flag_ = dtype_;
    tblob_.SetDLTensor(ptr_->shandle.ctx.dev_mask(), ptr_->shandle.ctx.dev_id);
#if MKL_EXPERIMENTAL == 1
    tblob_.Mkl_mem_ = Mkl_mem_;
#endif
  }

#if MKL_EXPERIMENTAL == 1
  std::shared_ptr<MKLMemHolder> Mkl_mem_;
#endif

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

size_t num_aux_data(NDArrayStorageType stype);

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

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

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 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_