Function ndarray_from_tensor

imperative/python/src/helper.cpp:599–651

}

PyObject* ndarray_from_tensor(const HostTensorND& val, ShareType share_type) {
    if (!val.layout().is_contiguous() && !val.shape().is_empty()) {
        mgb_assert(share_type != ShareType::MUST_SHARE);
        HostTensorND contig;
        contig.copy_from(val);
        return ndarray_from_tensor(contig, ShareType::TRY_SHARE);
    }
    PYTHON_GIL;
    npy_intp dims[TensorLayout::MAX_NDIM];
    for (size_t i = 0; i < val.layout().ndim; ++i)
        dims[i] = val.shape()[i];
    PyObject* ret = nullptr;

    auto alloc_new_ret = [&]() {
        mgb_assert(!ret);
        ret = PyArray_NewFromDescr(
                &PyArray_Type, dtype_mgb2np_descr(val.dtype()).release(),
                val.layout().ndim, dims, nullptr, nullptr, 0, nullptr);
        mgb_assert(ret, "failed to allocate array");
        mgb_assert(PyArray_Check(ret));
        return PyArray_DATA(reinterpret_cast<PyArrayObject*>(ret));
    };
    if (val.dtype().is_low_bit()) {
        mgb_assert(
                share_type != ShareType::MUST_SHARE,
                "can not share memory for lowbit dtype");
        const auto& layout = val.layout();
        if (layout.format.is_lowbit_aligned()) {
            lowbit_memcpy_aligned2byte(alloc_new_ret(), val.raw_ptr(), val.layout());
        } else {
            lowbit_memcpy_compact2byte(
                    val.dtype(), alloc_new_ret(), val.raw_ptr(),
                    val.layout().total_nr_elems());
        }
    } else if (share_type == ShareType::MUST_UNSHARE) {
        memcpy(alloc_new_ret(), val.raw_ptr(), val.layout().span().dist_byte());
    } else {
        // share data
        ret = PyArray_NewFromDescr(
                &PyArray_Type, dtype_mgb2np_descr(val.dtype()).release(),
                val.layout().ndim, dims, nullptr, const_cast<dt_byte*>(val.raw_ptr()),
                0, nullptr);
        mgb_assert(ret, "failed to alloc ndarray");
        auto capsule = PyCapsule_New(
                HostTensorNDRefHolder::alloc(val), "HostTensorND",
                ndarray_shared_from_tensor_py_capsule_dtor);
        mgb_assert(capsule, "failed to create PyCapsule");
        auto err =
                PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(ret), capsule);
        mgb_assert(!err);
    }
    return ret;
}

HostTensorND np2tensor(PyObject* obj, const Meth& meth, DType dtype) {
    auto ret_full = np2tensor_try_borrow(obj, meth, dtype);