Function _expand_dims_cpp

imperative/python/src/tensor_utils.cpp:1539–1573

}

py::object _expand_dims_cpp(py::handle inp_hdl, py::handle axis_hdl) {
    std::vector<int32_t> axis = list2vector(axis_hdl);
    bool unknown_ndim = true;
    size_t ndim = axis.size();
    if (auto p = TensorWrapper::try_cast(inp_hdl.ptr())) {
        auto&& shape = p->m_tensor->shape();
        if (shape) {
            unknown_ndim = false;
            ndim += shape->ndim;
        }
    } else {
        auto&& inp_ndim = get_ndim_safe(inp_hdl);
        ndim += inp_ndim.first;
        unknown_ndim &= !inp_ndim.second;
    }
    for (size_t i = 0; i < axis.size(); ++i) {
        if (axis[i] < 0) {
            if (unknown_ndim) {
                throw py::index_error(
                        "Does not support negative index when tensor's ndim is "
                        "unknown");
            }
            axis[i] += static_cast<int32_t>(ndim);
        }
    }
    if (!axis.size()) {
        throw py::index_error("axis could not be empty");
    }
    std::sort(axis.begin(), axis.end());
    std::shared_ptr<OpDef> op = AddAxis::make(axis = axis);
    py::object Op = py::cast(op);
    PyObject* p[2] = {Op.ptr(), inp_hdl.ptr()};
    py::tuple ret = py::reinterpret_steal<py::object>(py_apply(NULL, p, 2));
    return ret[0];
}

py::object _squeeze_cpp(py::handle inp_hdl, py::handle axis_hdl) {
    std::vector<int32_t> axis;