Function init_ops

imperative/python/src/ops.cpp:541–660

}  // namespace PYBIND11_NAMESPACE

void init_ops(py::module m) {
    _init_py_op_def(m);
    _init_py_op_base(m);
    INIT_ALL_OP(m)

    m.def("new_rng_handle", &rng::new_handle);
    m.def(
            "delete_rng_handle",
            [](size_t handle) {
                if (mgb::imperative::python::interpreter_for_py->check_available()) {
                    mgb::imperative::python::interpreter_for_py->sync();
                }
                mgb::CompNode::sync_all();
                mgb::CompNode::foreach ([](mgb::CompNode cn) {
                    auto err = cn.check_async_error();
                    mgb_assert(!err, "%s", err->what());
                });
                py_task_q.wait_all_task_finish();
                rng::delete_handle(handle);
            },
            py::call_guard<py::gil_scoped_release>());
    m.def("set_global_rng_seed", [](uint64_t seed) -> void {
        mgb_assert(
                python::interpreter_for_py->check_available(),
                "set global random seed failed since imperative interpreter has been "
                "destroyed");
        python::interpreter_for_py->sync();
        mgb::CompNode::sync_all();
        rng::set_global_rng_seed(seed);
    });
    m.def("get_global_rng_seed", &rng::get_global_rng_seed);
    m.def("get_rng_handle_compnode", &rng::get_rng_handle_compnode);

    struct PySubgraphBuilder {
        explicit PySubgraphBuilder(std::string name) : name{name} {}
        std::string name;
        Subgraph graph;
        mgb::SmallVector<bool> output_grad_mask;
        Subgraph::var_t next_var = 1;
        std::shared_ptr<mgb::Hashable> key = nullptr;

        std::shared_ptr<OpDef> build() {
            if (key == nullptr) {
                key = std::make_shared<UniqueKey>();
            }
            return SubgraphOp::make(
                    name, std::make_shared<Subgraph>(graph), output_grad_mask, key);
        }
    };

    py::class_<PySubgraphBuilder>(m, "SubgraphBuilder")
            .def(py::init<std::string>())
            .def(py::init<PySubgraphBuilder>())
            .def("input",
                 [](PySubgraphBuilder& self) {
                     mgb_assert(self.key == nullptr);
                     auto var = self.next_var++;
                     self.graph.inputs.push_back(var);