pass_functions.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 NNVM_PASS_FUNCTIONS_H_
#define NNVM_PASS_FUNCTIONS_H_

#include <string>
#include <memory>
#include <vector>
#include <utility>
#include "base.h"
#include "pass.h"
#include "graph_attr_types.h"

namespace nnvm {
namespace pass {

inline Graph LoadJSON(const std::string& json_str) {
  Graph ret;
  ret.attrs["json"] = std::make_shared<any>(json_str);
  return ApplyPass(ret, "LoadJSON");
}

inline std::string SaveJSON(Graph graph) {
  Graph ret = ApplyPass(std::move(graph), "SaveJSON");
  return ret.GetAttr<std::string>("json");
}


inline std::string PrintGraphIR(Graph graph) {
  Graph ret = ApplyPass(std::move(graph), "PrintGraphIR");
  return ret.GetAttr<std::string>("graphir");
}

inline Graph OrderMutation(Graph src) {
  return ApplyPass(std::move(src), "OrderMutation");
}

inline Graph InferShape(Graph graph,
                        ShapeVector shape_inputs,
                        std::string shape_attr_key = "") {
  if (shape_inputs.size() != 0) {
    graph.attrs["shape_inputs"] = std::make_shared<any>(std::move(shape_inputs));
  }
  if (shape_attr_key.length() != 0) {
    graph.attrs["shape_attr_key"] = std::make_shared<any>(std::move(shape_attr_key));
  }
  return ApplyPass(std::move(graph), "InferShape");
}

inline Graph InferType(Graph graph,
                       DTypeVector dtype_inputs,
                       std::string dtype_attr_key = "") {
  if (dtype_inputs.size() != 0) {
    graph.attrs["dtype_inputs"] = std::make_shared<any>(std::move(dtype_inputs));
  }
  if (dtype_attr_key.length() != 0) {
    graph.attrs["dtype_attr_key"] = std::make_shared<any>(std::move(dtype_attr_key));
  }
  return ApplyPass(std::move(graph), "InferType");
}

inline Graph PlaceDevice(Graph graph,
                         std::string device_group_attr_key,
                         DeviceAssignMap device_assign_map,
                         std::string device_copy_op) {
  graph.attrs["device_group_attr_key"] = std::make_shared<any>(std::move(device_group_attr_key));
  graph.attrs["device_assign_map"] = std::make_shared<any>(std::move(device_assign_map));
  graph.attrs["device_copy_op"] = std::make_shared<any>(std::move(device_copy_op));
  return ApplyPass(std::move(graph), "PlaceDevice");
}

inline Graph Gradient(
    Graph graph,
    std::vector<NodeEntry> ys,
    std::vector<NodeEntry> xs,
    std::vector<NodeEntry> ys_out_grad,
    std::function<NodeEntry(std::vector<NodeEntry>&& inputs)> aggregate_fun = nullptr,
    std::function<int(const Node& node)> mirror_fun = nullptr,
    std::function<NodeEntry(const NodeEntry& src, const NodeEntry &like)>
    attr_hint_fun = nullptr,
    std::vector<const Op*> zero_ops = std::vector<const Op*>(),
    std::string copy_op_str = std::string()) {
  graph.attrs["grad_ys"] = std::make_shared<any>(std::move(ys));

  graph.attrs["grad_xs"] = std::make_shared<any>(std::move(xs));
  graph.attrs["grad_ys_out_grad"] = std::make_shared<any>(std::move(ys_out_grad));
  if (aggregate_fun != nullptr) {
    graph.attrs["grad_aggregate_fun"] = std::make_shared<any>(aggregate_fun);
  }

  if (mirror_fun != nullptr) {
    graph.attrs["grad_mirror_fun"] = std::make_shared<any>(mirror_fun);
  }

  if (attr_hint_fun != nullptr) {
    graph.attrs["attr_hint_fun"] = std::make_shared<any>(attr_hint_fun);
  }

  if (zero_ops.size()) {
    graph.attrs["zero_ops"] = std::make_shared<any>(std::move(zero_ops));
  }

  if (copy_op_str != std::string()) {
      graph.attrs["copy_op"] = std::make_shared<any>(std::move(copy_op_str));
  }

  return ApplyPass(std::move(graph), "Gradient");
}

}  // namespace pass
}  // namespace nnvm
#endif  // NNVM_PASS_FUNCTIONS_H_