From 02ec955308f915cf4044473c8015fc521e3408cd Mon Sep 17 00:00:00 2001 From: Kaiqi Yan Date: Wed, 8 Jul 2026 07:12:01 +0000 Subject: [PATCH] handle empty tensor Signed-off-by: Kaiqi Yan --- libs/qec/python/bindings/py_code.cpp | 38 ++++++++++++++++++------- libs/qec/python/bindings/type_casters.h | 15 ++++++++-- libs/qec/python/tests/test_code.py | 22 ++++++++++++++ 3 files changed, 62 insertions(+), 13 deletions(-) diff --git a/libs/qec/python/bindings/py_code.cpp b/libs/qec/python/bindings/py_code.cpp index 8e3c493d..2d0ba7e2 100644 --- a/libs/qec/python/bindings/py_code.cpp +++ b/libs/qec/python/bindings/py_code.cpp @@ -227,6 +227,21 @@ static nb::object get_python_kernel_or_throw(const code &self, operation op) { return it->second; } + +template +static auto copyCodeMatrixToPyArray(const cudaqx::tensor &tensor, + std::size_t emptyCols) { + if (tensor.rank() == 0 && tensor.size() == 0) { + T *data_copy = new T[0]; + size_t arr_shape[] = {0, emptyCols}; + return nb::ndarray( + data_copy, 2, arr_shape, nb::capsule(data_copy, [](void *p) noexcept { + delete[] static_cast(p); + })); + } + + return cudaq::python::copyCUDAQXTensorToPyArray(tensor); +} } // namespace // Registry to store code factory functions @@ -364,42 +379,43 @@ void bindCode(nb::module_ &mod) { .def( "get_parity", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray(code.get_parity()); + return copyCodeMatrixToPyArray(code.get_parity(), + 2 * code.get_num_data_qubits()); }, "Get the parity check matrix of the code") .def( "get_parity_x", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray( - code.get_parity_x()); + return copyCodeMatrixToPyArray(code.get_parity_x(), + code.get_num_data_qubits()); }, "Get the X-type parity check matrix of the code") .def( "get_parity_z", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray( - code.get_parity_z()); + return copyCodeMatrixToPyArray(code.get_parity_z(), + code.get_num_data_qubits()); }, "Get the Z-type parity check matrix of the code") .def( "get_pauli_observables_matrix", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray( - code.get_pauli_observables_matrix()); + return copyCodeMatrixToPyArray(code.get_pauli_observables_matrix(), + 2 * code.get_num_data_qubits()); }, "Get a matrix of the Pauli observables of the code") .def( "get_observables_x", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray( - code.get_observables_x()); + return copyCodeMatrixToPyArray(code.get_observables_x(), + code.get_num_data_qubits()); }, "Get the Pauli X observables of the code") .def( "get_observables_z", [](code &code) { - return cudaq::python::copyCUDAQXTensorToPyArray( - code.get_observables_z()); + return copyCodeMatrixToPyArray(code.get_observables_z(), + code.get_num_data_qubits()); }, "Get the Pauli Z observables of the code") .def("get_stabilizers", &code::get_stabilizers, diff --git a/libs/qec/python/bindings/type_casters.h b/libs/qec/python/bindings/type_casters.h index cc122767..72c98df4 100644 --- a/libs/qec/python/bindings/type_casters.h +++ b/libs/qec/python/bindings/type_casters.h @@ -8,6 +8,7 @@ #pragma once #include +#include #include "common/ObserveResult.h" #include "cuda-qx/core/heterogeneous_map.h" @@ -256,13 +257,18 @@ namespace python { template auto copyCUDAQXTensorToPyArray(const cudaqx::tensor &tensor) { auto shape = tensor.shape(); + if (shape.size() != 2) + throw std::runtime_error( + "Expected a rank-2 cudaqx::tensor for NumPy conversion."); + auto rows = shape[0]; auto cols = shape[1]; size_t total_size = rows * cols; // Allocate new memory and copy the data T *data_copy = new T[total_size]; - std::memcpy(data_copy, tensor.data(), total_size * sizeof(T)); + if (total_size > 0) + std::memcpy(data_copy, tensor.data(), total_size * sizeof(T)); size_t arr_shape[] = {rows, cols}; return nb::ndarray(data_copy, 2, arr_shape, @@ -274,12 +280,17 @@ auto copyCUDAQXTensorToPyArray(const cudaqx::tensor &tensor) { template auto copy1DCUDAQXTensorToPyArray(const cudaqx::tensor &tensor) { auto shape = tensor.shape(); + if (shape.size() != 1) + throw std::runtime_error( + "Expected a rank-1 cudaqx::tensor for NumPy conversion."); + auto rows = shape[0]; size_t total_size = rows; // Allocate new memory and copy the data T *data_copy = new T[total_size]; - std::memcpy(data_copy, tensor.data(), total_size * sizeof(T)); + if (total_size > 0) + std::memcpy(data_copy, tensor.data(), total_size * sizeof(T)); size_t arr_shape[] = {rows}; return nb::ndarray(data_copy, 1, arr_shape, diff --git a/libs/qec/python/tests/test_code.py b/libs/qec/python/tests/test_code.py index 3b049236..f6f158e1 100644 --- a/libs/qec/python/tests/test_code.py +++ b/libs/qec/python/tests/test_code.py @@ -39,6 +39,28 @@ def test_code_parity_matrices(): assert parity_z.shape == (3, 7) +def test_repetition_empty_x_matrices_preserve_rank(): + repetition = qec.get_code("repetition", distance=3) + + # Repetition is Z-only; empty X-side results must still be matrices with + # one column per data qubit so Python callers can inspect their shape. + parity_x = repetition.get_parity_x() + assert isinstance(parity_x, np.ndarray) + assert parity_x.dtype == np.uint8 + assert parity_x.shape == (0, 3) + + observables_x = repetition.get_observables_x() + assert isinstance(observables_x, np.ndarray) + assert observables_x.dtype == np.uint8 + assert observables_x.shape == (0, 3) + + parity_z = repetition.get_parity_z() + assert parity_z.shape == (2, 3) + + observables_z = repetition.get_observables_z() + assert observables_z.shape == (1, 3) + + def test_code_stabilizers(): steane = qec.get_code("steane") stabilizers = steane.get_stabilizers()