Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 27 additions & 11 deletions libs/qec/python/bindings/py_code.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,21 @@ static nb::object get_python_kernel_or_throw(const code &self, operation op) {

return it->second;
}

template <typename T>
static auto copyCodeMatrixToPyArray(const cudaqx::tensor<T> &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<nb::numpy, T>(
data_copy, 2, arr_shape, nb::capsule(data_copy, [](void *p) noexcept {
delete[] static_cast<T *>(p);
}));
}

return cudaq::python::copyCUDAQXTensorToPyArray(tensor);
}
} // namespace

// Registry to store code factory functions
Expand Down Expand Up @@ -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,
Expand Down
15 changes: 13 additions & 2 deletions libs/qec/python/bindings/type_casters.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#pragma once

#include <cstring>
#include <stdexcept>

#include "common/ObserveResult.h"
#include "cuda-qx/core/heterogeneous_map.h"
Expand Down Expand Up @@ -256,13 +257,18 @@ namespace python {
template <typename T>
auto copyCUDAQXTensorToPyArray(const cudaqx::tensor<T> &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<nb::numpy, T>(data_copy, 2, arr_shape,
Expand All @@ -274,12 +280,17 @@ auto copyCUDAQXTensorToPyArray(const cudaqx::tensor<T> &tensor) {
template <typename T>
auto copy1DCUDAQXTensorToPyArray(const cudaqx::tensor<T> &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<nb::numpy, T>(data_copy, 1, arr_shape,
Expand Down
22 changes: 22 additions & 0 deletions libs/qec/python/tests/test_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading