diff --git a/imperative/python/src/helper.cpp b/imperative/python/src/helper.cpp index 98d78ea06549dd84be2d5d4ad9a0d9ad0b753edf..7f38403b504c6e09753c551867f67fd0a722f07c 100644 --- a/imperative/python/src/helper.cpp +++ b/imperative/python/src/helper.cpp @@ -192,12 +192,16 @@ std::unique_ptr dtype_mgb2np_descr(DType dty const std::vector>& data) { PyObject* metadata = PyDict_New(); PyObject* mgb_dtype_metadata = PyDict_New(); - PyDict_SetItemString( - mgb_dtype_metadata, "name", PyUnicode_FromString(name)); + PyObject* py_name = PyUnicode_FromString(name); + + PyDict_SetItemString(mgb_dtype_metadata, "name", py_name); + Py_DECREF(py_name); + for (const auto& d : data) { PyDict_SetItemString(mgb_dtype_metadata, d.first, d.second); } PyDict_SetItemString(metadata, "mgb_dtype", mgb_dtype_metadata); + Py_DECREF(mgb_dtype_metadata); return metadata; }; if (dtype.has_param()) { @@ -206,51 +210,62 @@ std::unique_ptr dtype_mgb2np_descr(DType dty case DTypeEnum::QuantizedS1: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_INT8); + PyObject* scale = PyFloat_FromDouble(param.scale); type_descr->metadata = build_mgb_dtype_dict( - DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}}); + DTypeTrait::name, {{"scale", scale}}); + Py_DECREF(scale); break; } case DTypeEnum::Quantized4Asymm: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_UINT8); + PyObject* scale = PyFloat_FromDouble(param.scale); + PyObject* zero_point = PyLong_FromLong(param.zero_point); type_descr->metadata = build_mgb_dtype_dict( DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}, - {"zero_point", PyLong_FromLong(param.zero_point)}}); + {{"scale", scale}, {"zero_point", zero_point}}); + Py_DECREF(scale); + Py_DECREF(zero_point); break; } case DTypeEnum::QuantizedS4: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_INT8); + PyObject* scale = PyFloat_FromDouble(param.scale); type_descr->metadata = build_mgb_dtype_dict( - DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}}); + DTypeTrait::name, {{"scale", scale}}); + Py_DECREF(scale); break; } case DTypeEnum::Quantized8Asymm: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_UINT8); + PyObject* scale = PyFloat_FromDouble(param.scale); + PyObject* zero_point = PyLong_FromLong(param.zero_point); type_descr->metadata = build_mgb_dtype_dict( DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}, - {"zero_point", PyLong_FromLong(param.zero_point)}}); + {{"scale", scale}, {"zero_point", zero_point}}); + Py_DECREF(scale); + Py_DECREF(zero_point); break; } case DTypeEnum::QuantizedS8: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_INT8); - type_descr->metadata = build_mgb_dtype_dict( - DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}}); + PyObject* scale = PyFloat_FromDouble(param.scale); + auto metadata = build_mgb_dtype_dict( + DTypeTrait::name, {{"scale", scale}}); + type_descr->metadata = metadata; + Py_DECREF(scale); break; } case DTypeEnum::QuantizedS32: { auto& param = dtype.param(); type_descr = PyArray_DescrNewFromType(NPY_INT32); + PyObject* scale = PyFloat_FromDouble(param.scale); type_descr->metadata = build_mgb_dtype_dict( - DTypeTrait::name, - {{"scale", PyFloat_FromDouble(param.scale)}}); + DTypeTrait::name, {{"scale", scale}}); + Py_DECREF(scale); break; } default: diff --git a/imperative/python/src/tensor_utils.cpp b/imperative/python/src/tensor_utils.cpp index 906615f5d3c42c65276102ea4962e652b7ae213f..19a2f3f7d1e211ce1a2a0135634ae702e6995882 100644 --- a/imperative/python/src/tensor_utils.cpp +++ b/imperative/python/src/tensor_utils.cpp @@ -422,6 +422,7 @@ mgb::DType _get_dtype(py::handle tensor) { py::object _astype_cpp(py::handle tensor, py::handle dtype_hdl) { PyArray_Descr* descr; + py::object ret; if (!PyArray_DescrConverter(dtype_hdl.ptr(), &descr)) { throw py::value_error(ssprintf( "can not convert to numpy.dtype from %s", @@ -432,11 +433,13 @@ py::object _astype_cpp(py::handle tensor, py::handle dtype_hdl) { std::shared_ptr op = TypeCvt::make(npy::dtype_np2mgb_descr(descr)); py::object Op = py::cast(op); PyObject* p[2] = {Op.ptr(), tensor.ptr()}; - py::tuple ret = py::reinterpret_steal(py_apply(NULL, p, 2)); - return ret[0]; + py::tuple apply_res = py::reinterpret_steal(py_apply(NULL, p, 2)); + ret = apply_res[0]; } else { - return py::reinterpret_borrow(tensor); + ret = py::reinterpret_borrow(tensor); } + Py_DECREF(descr); + return ret; } py::object _convert_single_value_cpp(