operator.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_CPP_OPERATOR_H_
#define MXNET_CPP_OPERATOR_H_

#include <map>
#include <string>
#include <vector>
#include "mxnet-cpp/base.h"
#include "mxnet-cpp/op_map.h"
#include "mxnet-cpp/symbol.h"

namespace mxnet {
namespace cpp {
class Mxnet;
class Operator {
 public:
  explicit Operator(const std::string &operator_name);
  Operator &operator=(const Operator &rhs);
  template <typename T>
  Operator &SetParam(const std::string &name, const T &value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;

    params_[name] = value_str;
    return *this;
  }
  template <typename T>
  Operator &SetParam(int pos, const T &value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;

    params_[arg_names_[pos]] = value_str;
    return *this;
  }
  Operator &SetInput(const std::string &name, const Symbol &symbol);
  template<int N = 0>
  void PushInput(const Symbol &symbol) {
    input_symbols_.push_back(symbol.GetHandle());
  }
  Operator &operator()() { return *this; }
  Operator &operator()(const Symbol &symbol) {
    input_symbols_.push_back(symbol.GetHandle());
    return *this;
  }
  Operator &operator()(const std::vector<Symbol> &symbols) {
    for (auto &s : symbols) {
      input_symbols_.push_back(s.GetHandle());
    }
    return *this;
  }
  Symbol CreateSymbol(const std::string &name = "");

  Operator &SetInput(const std::string &name, const NDArray &ndarray);
  template<int N = 0>
  Operator &PushInput(const NDArray &ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  template <class T, class... Args, int N = 0>
  Operator &PushInput(const T &t, Args... args) {
    SetParam(N, t);
    PushInput<Args..., N+1>(args...);
    return *this;
  }
  template <class T, int N = 0>
  Operator &PushInput(const T &t) {
    SetParam(N, t);
    return *this;
  }
  Operator &operator()(const NDArray &ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  Operator &operator()(const std::vector<NDArray> &ndarrays) {
    for (auto &s : ndarrays) {
      input_ndarrays_.push_back(s.GetHandle());
    }
    return *this;
  }
  template <typename... Args>
  Operator &operator()(Args... args) {
    PushInput(args...);
    return *this;
  }
  std::vector<NDArray> Invoke();
  void Invoke(NDArray &output);
  void Invoke(std::vector<NDArray> &outputs);

 private:
  std::map<std::string, std::string> params_desc_;
  bool variable_params_ = false;
  std::map<std::string, std::string> params_;
  std::vector<SymbolHandle> input_symbols_;
  std::vector<NDArrayHandle> input_ndarrays_;
  std::vector<std::string> input_keys_;
  std::vector<std::string> arg_names_;
  AtomicSymbolCreator handle_;
  static OpMap*& op_map();
};
}  // namespace cpp
}  // namespace mxnet

#endif  // MXNET_CPP_OPERATOR_H_