MegEngine/imperative/python/src/imperative_rt.cpp

/**
 * \file imperative/python/src/imperative_rt.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "./imperative_rt.h"

#include <future>
#include <variant>
#include <unordered_map>
#include <pybind11/numpy.h>
#include <pybind11/operators.h>

#include "megbrain/imperative.h"
#include "megbrain/imperative/interpreter.h"
#include "megbrain/imperative/ops/opr_attr.h"
#include "./helper.h"
#include "./common.h"

namespace py = pybind11;

using namespace mgb;
using namespace imperative;
using namespace interpreter;


void init_imperative_rt(py::module m) {
    auto make_backward_graph = [](
            const OpDef& def,
            const SmallVector<LogicalTensorDesc>& inputs,
            const SmallVector<bool>& input_requires_grad,
            const SmallVector<bool>& output_has_grad){
        auto result = OpDef::make_backward_graph(def, inputs, input_requires_grad, output_has_grad);
        return std::make_tuple("backward_graph", result.save_for_backward, result.input_has_grad);
    };
    m.def("make_backward_graph", make_backward_graph);
}