fix: To fix Python QEC conversion for empty repetition-code X matrices (#660)

This PR tries to fix [[B]
6428084](https://nvbugspro.nvidia.com/bug/6428084) .
It also adds one test to expose this bug and verify the fix.
Thanks for looking this PR.

In our suggested fix, we keep the existing C++ `cudaqx::tensor`
semantics unchanged: an empty directional parity/observable tensor may
remain rank 0.
In the Python binding layer, we add a small wrapper for matrix-returning
`Code` methods that detects this valid rank-0 empty sentinel and
converts it to an empty 2-D NumPy matrix with the correct width, e.g.
`(0, code.get_num_data_qubits())`.

It also adds rank validation in the generic tensor-to-NumPy helpers
before reading `shape[0]` / `shape[1]` and skips `memcpy` for zero-sized
tensors, so invalid-rank conversions fail cleanly instead of causing
undefined behaviour.

Signed-off-by: Kaiqi Yan <kaiqiy@nvidia.com>

Author: kaiqiy-nv

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 <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
@@ -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,

libs/qec/python/bindings/type_casters.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <cstring>
+#include <stdexcept>
 
 #include "common/ObserveResult.h"
 #include "cuda-qx/core/heterogeneous_map.h"
@@ -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,
@@ -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,

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()