This is an automated email from the ASF dual-hosted git repository.
guanmingchiu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/mahout.git
The following commit(s) were added to refs/heads/main by this push:
new 11db508df [QDP] Refactor `encode()` method into helper functions with
tests (#814)
11db508df is described below
commit 11db508dfdfc5deac48f4f0250f421dcbc8ae92c
Author: Vic Wen <[email protected]>
AuthorDate: Fri Jan 23 23:39:55 2026 +0800
[QDP] Refactor `encode()` method into helper functions with tests (#814)
* refactor: Refactor encoding methods for NumPy and PyTorch tensors
- Introduced dedicated methods `encode_from_numpy` and
`encode_from_pytorch` to streamline the encoding process for NumPy arrays and
PyTorch tensors, respectively.
- Improved error handling for unsupported shapes in both encoding methods.
- Simplified the main encoding logic by delegating to these new methods.
* chore: remove unused import
* fix(python): replace PyReadonlyArrayDyn with PyReadonlyArray1/2
Replace deprecated PyReadonlyArrayDyn with dimension-specific types
PyReadonlyArray1 and PyReadonlyArray2 for better type safety and
compatibility with newer pyo3-numpy versions.
* refactor(tests): refactor `test_bindings.py`
* feat(validation): add shape validation for arrays and tensors
* test: add 3D tensor shape validation testing
* Remove commented-out section for IQP Encoding Tests in `test_bindings.py`
---
qdp/qdp-python/src/lib.rs | 544 ++++++++++++++++++++++++++-----------------
testing/qdp/test_bindings.py | 348 ++++++++++++++-------------
2 files changed, 507 insertions(+), 385 deletions(-)
diff --git a/qdp/qdp-python/src/lib.rs b/qdp/qdp-python/src/lib.rs
index 016ee1259..bcf03c129 100644
--- a/qdp/qdp-python/src/lib.rs
+++ b/qdp/qdp-python/src/lib.rs
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyReadonlyArrayDyn,
PyUntypedArrayMethods};
+use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods};
use pyo3::exceptions::PyRuntimeError;
use pyo3::ffi;
use pyo3::prelude::*;
@@ -178,6 +178,32 @@ fn is_cuda_tensor(tensor: &Bound<'_, PyAny>) ->
PyResult<bool> {
Ok(device_type == "cuda")
}
+/// Validate array/tensor shape (must be 1D or 2D)
+///
+/// Args:
+/// ndim: Number of dimensions
+/// context: Context string for error message (e.g., "array", "tensor",
"CUDA tensor")
+///
+/// Returns:
+/// Ok(()) if shape is valid (1D or 2D), otherwise returns an error
+fn validate_shape(ndim: usize, context: &str) -> PyResult<()> {
+ match ndim {
+ 1 | 2 => Ok(()),
+ _ => {
+ let item_type = if context.contains("array") {
+ "array"
+ } else {
+ "tensor"
+ };
+ Err(PyRuntimeError::new_err(format!(
+ "Unsupported {} shape: {}D. Expected 1D {} for single sample \
+ encoding or 2D {} (batch_size, features) for batch encoding.",
+ context, ndim, item_type, item_type
+ )))
+ }
+ }
+}
+
/// Get the CUDA device index from a PyTorch tensor
fn get_tensor_device_id(tensor: &Bound<'_, PyAny>) -> PyResult<i32> {
let device = tensor.getattr("device")?;
@@ -238,17 +264,43 @@ fn validate_cuda_tensor_for_encoding(
Ok(())
}
-/// CUDA tensor information extracted directly from PyTorch tensor
-struct CudaTensorInfo {
+/// DLPack tensor information extracted from a PyCapsule
+///
+/// This struct owns the DLManagedTensor pointer and ensures proper cleanup
+/// via the DLPack deleter when dropped (RAII pattern).
+struct DLPackTensorInfo {
+ /// Raw DLManagedTensor pointer from PyTorch DLPack capsule
+ /// This is owned by this struct and will be freed via deleter on drop
+ managed_ptr: *mut DLManagedTensor,
+ /// Data pointer inside dl_tensor (GPU memory, owned by managed_ptr)
data_ptr: *const f64,
shape: Vec<i64>,
+ /// CUDA device ID from DLPack metadata.
+ /// Currently unused but kept for potential future device validation or
multi-GPU support.
+ #[allow(dead_code)]
+ device_id: i32,
+}
+
+impl Drop for DLPackTensorInfo {
+ fn drop(&mut self) {
+ unsafe {
+ if !self.managed_ptr.is_null() {
+ // Per DLPack protocol: consumer must call deleter exactly once
+ if let Some(deleter) = (*self.managed_ptr).deleter {
+ deleter(self.managed_ptr);
+ }
+ // Prevent double-free
+ self.managed_ptr = std::ptr::null_mut();
+ }
+ }
+ }
}
-/// Extract GPU pointer directly from PyTorch CUDA tensor
+/// Extract GPU pointer from PyTorch tensor's __dlpack__() capsule
///
-/// Uses PyTorch's `data_ptr()` and `shape` APIs directly instead of DLPack
protocol.
-/// This avoids the DLPack capsule lifecycle complexity and potential memory
leaks
-/// from the capsule renaming pattern.
+/// Uses the DLPack protocol to obtain a zero-copy view of the tensor's GPU
memory.
+/// The returned `DLPackTensorInfo` owns the DLManagedTensor and will
automatically
+/// call the deleter when dropped, ensuring proper resource cleanup.
///
/// # Safety
/// The returned `data_ptr` points to GPU memory owned by the source tensor.
@@ -256,19 +308,59 @@ struct CudaTensorInfo {
/// for the entire duration that `data_ptr` is in use. Python's GIL ensures
/// the tensor won't be garbage collected during `encode()`, but the caller
/// must not deallocate or resize the tensor while encoding is in progress.
-fn extract_cuda_tensor_info(tensor: &Bound<'_, PyAny>) ->
PyResult<CudaTensorInfo> {
- // Get GPU pointer directly via tensor.data_ptr()
- let data_ptr_int: isize = tensor.call_method0("data_ptr")?.extract()?;
- if data_ptr_int == 0 {
- return Err(PyRuntimeError::new_err("CUDA tensor has null data
pointer"));
- }
- let data_ptr = data_ptr_int as *const f64;
+fn extract_dlpack_tensor(_py: Python<'_>, tensor: &Bound<'_, PyAny>) ->
PyResult<DLPackTensorInfo> {
+ // Call tensor.__dlpack__() to get PyCapsule
+ // Note: PyTorch's __dlpack__ uses the default stream when called without
arguments
+ let capsule = tensor.call_method0("__dlpack__")?;
+
+ // Extract the DLManagedTensor pointer from the capsule
+ const DLTENSOR_NAME: &[u8] = b"dltensor\0";
+
+ unsafe {
+ let capsule_ptr = capsule.as_ptr();
+ let managed_ptr =
+ ffi::PyCapsule_GetPointer(capsule_ptr, DLTENSOR_NAME.as_ptr() as
*const i8)
+ as *mut DLManagedTensor;
+
+ if managed_ptr.is_null() {
+ return Err(PyRuntimeError::new_err(
+ "Failed to extract DLManagedTensor from PyCapsule",
+ ));
+ }
- // Get shape directly via tensor.shape
- let shape_obj = tensor.getattr("shape")?;
- let shape: Vec<i64> = shape_obj.extract()?;
+ let dl_tensor = &(*managed_ptr).dl_tensor;
+
+ // Extract data pointer with null check
+ if dl_tensor.data.is_null() {
+ return Err(PyRuntimeError::new_err(
+ "DLPack tensor has null data pointer",
+ ));
+ }
+ let data_ptr = dl_tensor.data as *const f64;
+
+ // Extract shape
+ let ndim = dl_tensor.ndim as usize;
+ let shape = if ndim > 0 && !dl_tensor.shape.is_null() {
+ std::slice::from_raw_parts(dl_tensor.shape, ndim).to_vec()
+ } else {
+ vec![]
+ };
- Ok(CudaTensorInfo { data_ptr, shape })
+ // Extract device_id
+ let device_id = dl_tensor.device.device_id;
+
+ // Rename the capsule to "used_dltensor" as per DLPack protocol
+ // This prevents PyTorch from trying to delete it when the capsule is
garbage collected
+ const USED_DLTENSOR_NAME: &[u8] = b"used_dltensor\0";
+ ffi::PyCapsule_SetName(capsule_ptr, USED_DLTENSOR_NAME.as_ptr() as
*const i8);
+
+ Ok(DLPackTensorInfo {
+ managed_ptr,
+ data_ptr,
+ shape,
+ device_id,
+ })
+ }
}
/// PyO3 wrapper for QdpEngine
@@ -358,237 +450,253 @@ impl QdpEngine {
// Check if it's a NumPy array
if data.hasattr("__array_interface__")? {
- // Get the array's ndim for shape validation
- let ndim: usize = data.getattr("ndim")?.extract()?;
-
- match ndim {
- 1 => {
- // 1D array: single sample encoding (zero-copy if already
contiguous)
- let array_1d =
data.extract::<PyReadonlyArray1<f64>>().map_err(|_| {
- PyRuntimeError::new_err(
- "Failed to extract 1D NumPy array. Ensure dtype is
float64.",
- )
- })?;
- let data_slice = array_1d.as_slice().map_err(|_| {
- PyRuntimeError::new_err("NumPy array must be
contiguous (C-order)")
- })?;
- let ptr = self
- .engine
- .encode(data_slice, num_qubits, encoding_method)
- .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
- return Ok(QuantumTensor {
- ptr,
- consumed: false,
- });
- }
- 2 => {
- // 2D array: batch encoding (zero-copy if already
contiguous)
- let array_2d =
data.extract::<PyReadonlyArray2<f64>>().map_err(|_| {
- PyRuntimeError::new_err(
- "Failed to extract 2D NumPy array. Ensure dtype is
float64.",
- )
- })?;
- let shape = array_2d.shape();
- let num_samples = shape[0];
- let sample_size = shape[1];
- let data_slice = array_2d.as_slice().map_err(|_| {
- PyRuntimeError::new_err("NumPy array must be
contiguous (C-order)")
- })?;
- let ptr = self
- .engine
- .encode_batch(
- data_slice,
- num_samples,
- sample_size,
- num_qubits,
- encoding_method,
- )
- .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
- return Ok(QuantumTensor {
- ptr,
- consumed: false,
- });
- }
- _ => {
- return Err(PyRuntimeError::new_err(format!(
- "Unsupported array shape: {}D. Expected 1D array for
single sample \
- encoding or 2D array (batch_size, features) for batch
encoding.",
- ndim
- )));
- }
- }
+ return self.encode_from_numpy(data, num_qubits, encoding_method);
}
// Check if it's a PyTorch tensor
if is_pytorch_tensor(data)? {
- // Check if it's a CUDA tensor - use zero-copy GPU encoding
- if is_cuda_tensor(data)? {
- // Validate CUDA tensor for direct GPU encoding
- validate_cuda_tensor_for_encoding(
- data,
- self.engine.device().ordinal(),
- encoding_method,
- )?;
-
- // Extract GPU pointer directly from PyTorch tensor
- let tensor_info = extract_cuda_tensor_info(data)?;
-
- let ndim: usize = data.call_method0("dim")?.extract()?;
-
- match ndim {
- 1 => {
- // 1D CUDA tensor: single sample encoding
- let input_len = tensor_info.shape[0] as usize;
- // SAFETY: tensor_info.data_ptr was obtained via
PyTorch's data_ptr() from a
- // valid CUDA tensor. The tensor remains alive during
this call
- // (held by Python's GIL), and we validated
dtype/contiguity/device above.
- let ptr = unsafe {
- self.engine
- .encode_from_gpu_ptr(
- tensor_info.data_ptr,
- input_len,
- num_qubits,
- encoding_method,
- )
- .map_err(|e| {
- PyRuntimeError::new_err(format!("Encoding
failed: {}", e))
- })?
- };
- return Ok(QuantumTensor {
- ptr,
- consumed: false,
- });
- }
- 2 => {
- // 2D CUDA tensor: batch encoding
- let num_samples = tensor_info.shape[0] as usize;
- let sample_size = tensor_info.shape[1] as usize;
- // SAFETY: Same as above - pointer from validated
PyTorch CUDA tensor
- let ptr = unsafe {
- self.engine
- .encode_batch_from_gpu_ptr(
- tensor_info.data_ptr,
- num_samples,
- sample_size,
- num_qubits,
- encoding_method,
- )
- .map_err(|e| {
- PyRuntimeError::new_err(format!("Encoding
failed: {}", e))
- })?
- };
- return Ok(QuantumTensor {
- ptr,
- consumed: false,
- });
- }
- _ => {
- return Err(PyRuntimeError::new_err(format!(
- "Unsupported CUDA tensor shape: {}D. Expected 1D
tensor for single \
- sample encoding or 2D tensor (batch_size,
features) for batch encoding.",
- ndim
- )));
- }
- }
- }
+ return self.encode_from_pytorch(data, num_qubits, encoding_method);
+ }
- // CPU tensor path (existing code)
- validate_tensor(data)?;
- // PERF: Avoid Tensor -> Python list -> Vec deep copies.
- //
- // For CPU tensors, `tensor.detach().numpy()` returns a NumPy view
that shares the same
- // underlying memory (zero-copy) when the tensor is C-contiguous.
We can then borrow a
- // `&[f64]` directly via pyo3-numpy.
- let ndim: usize = data.call_method0("dim")?.extract()?;
- let numpy_view = data
- .call_method0("detach")?
- .call_method0("numpy")
- .map_err(|_| {
+ // Fallback: try to extract as Vec<f64> (Python list)
+ self.encode_from_list(data, num_qubits, encoding_method)
+ }
+
+ /// Encode from NumPy array (1D or 2D)
+ fn encode_from_numpy(
+ &self,
+ data: &Bound<'_, PyAny>,
+ num_qubits: usize,
+ encoding_method: &str,
+ ) -> PyResult<QuantumTensor> {
+ let ndim: usize = data.getattr("ndim")?.extract()?;
+ validate_shape(ndim, "array")?;
+
+ match ndim {
+ 1 => {
+ // 1D array: single sample encoding (zero-copy if already
contiguous)
+ let array_1d =
data.extract::<PyReadonlyArray1<f64>>().map_err(|_| {
PyRuntimeError::new_err(
- "Failed to convert torch.Tensor to NumPy view. Ensure
the tensor is on CPU \
- and does not require grad (try: tensor =
tensor.detach().cpu())",
+ "Failed to extract 1D NumPy array. Ensure dtype is
float64.",
)
})?;
-
- let array = numpy_view
- .extract::<PyReadonlyArrayDyn<f64>>()
- .map_err(|_| {
+ let data_slice = array_1d.as_slice().map_err(|_| {
+ PyRuntimeError::new_err("NumPy array must be contiguous
(C-order)")
+ })?;
+ let ptr = self
+ .engine
+ .encode(data_slice, num_qubits, encoding_method)
+ .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ Ok(QuantumTensor {
+ ptr,
+ consumed: false,
+ })
+ }
+ 2 => {
+ // 2D array: batch encoding (zero-copy if already contiguous)
+ let array_2d =
data.extract::<PyReadonlyArray2<f64>>().map_err(|_| {
PyRuntimeError::new_err(
- "Failed to extract NumPy view as float64 array. Ensure
dtype is float64 \
- (try: tensor = tensor.to(torch.float64))",
+ "Failed to extract 2D NumPy array. Ensure dtype is
float64.",
)
})?;
+ let shape = array_2d.shape();
+ let num_samples = shape[0];
+ let sample_size = shape[1];
+ let data_slice = array_2d.as_slice().map_err(|_| {
+ PyRuntimeError::new_err("NumPy array must be contiguous
(C-order)")
+ })?;
+ let ptr = self
+ .engine
+ .encode_batch(
+ data_slice,
+ num_samples,
+ sample_size,
+ num_qubits,
+ encoding_method,
+ )
+ .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ Ok(QuantumTensor {
+ ptr,
+ consumed: false,
+ })
+ }
+ _ => unreachable!("validate_shape() should have caught invalid
ndim"),
+ }
+ }
- let data_slice = array.as_slice().map_err(|_| {
- PyRuntimeError::new_err(
- "Tensor must be contiguous (C-order) to get zero-copy
slice \
- (try: tensor = tensor.contiguous())",
- )
- })?;
+ /// Encode from PyTorch tensor (1D or 2D)
+ fn encode_from_pytorch(
+ &self,
+ data: &Bound<'_, PyAny>,
+ num_qubits: usize,
+ encoding_method: &str,
+ ) -> PyResult<QuantumTensor> {
+ // Check if it's a CUDA tensor - use zero-copy GPU encoding via DLPack
+ if is_cuda_tensor(data)? {
+ // Validate CUDA tensor for direct GPU encoding
+ validate_cuda_tensor_for_encoding(
+ data,
+ self.engine.device().ordinal(),
+ encoding_method,
+ )?;
+
+ // Extract GPU pointer via DLPack (RAII wrapper ensures deleter is
called)
+ let dlpack_info = extract_dlpack_tensor(data.py(), data)?;
+
+ let ndim: usize = data.call_method0("dim")?.extract()?;
+ validate_shape(ndim, "CUDA tensor")?;
match ndim {
1 => {
- // 1D tensor: single sample encoding
- let ptr = self
- .engine
- .encode(data_slice, num_qubits, encoding_method)
- .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ // 1D CUDA tensor: single sample encoding
+ let input_len = dlpack_info.shape[0] as usize;
+ // SAFETY: dlpack_info.data_ptr was validated via DLPack
protocol from a
+ // valid PyTorch CUDA tensor. The tensor remains alive
during this call
+ // (held by Python's GIL), and we validated
dtype/contiguity/device above.
+ // The DLPackTensorInfo RAII wrapper will call deleter
when dropped.
+ let ptr = unsafe {
+ self.engine
+ .encode_from_gpu_ptr(
+ dlpack_info.data_ptr,
+ input_len,
+ num_qubits,
+ encoding_method,
+ )
+ .map_err(|e| {
+ PyRuntimeError::new_err(format!("Encoding
failed: {}", e))
+ })?
+ };
return Ok(QuantumTensor {
ptr,
consumed: false,
});
}
2 => {
- // 2D tensor: batch encoding
- let shape = array.shape();
- if shape.len() != 2 {
- return Err(PyRuntimeError::new_err(format!(
- "Unsupported tensor shape: {}D. Expected 2D tensor
(batch_size, features).",
- shape.len()
- )));
- }
- let num_samples = shape[0];
- let sample_size = shape[1];
- let ptr = self
- .engine
- .encode_batch(
- data_slice,
- num_samples,
- sample_size,
- num_qubits,
- encoding_method,
- )
- .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ // 2D CUDA tensor: batch encoding
+ let num_samples = dlpack_info.shape[0] as usize;
+ let sample_size = dlpack_info.shape[1] as usize;
+ // SAFETY: Same as above - pointer from validated DLPack
tensor
+ let ptr = unsafe {
+ self.engine
+ .encode_batch_from_gpu_ptr(
+ dlpack_info.data_ptr,
+ num_samples,
+ sample_size,
+ num_qubits,
+ encoding_method,
+ )
+ .map_err(|e| {
+ PyRuntimeError::new_err(format!("Encoding
failed: {}", e))
+ })?
+ };
return Ok(QuantumTensor {
ptr,
consumed: false,
});
}
- _ => {
- return Err(PyRuntimeError::new_err(format!(
- "Unsupported tensor shape: {}D. Expected 1D tensor for
single sample \
- encoding or 2D tensor (batch_size, features) for
batch encoding.",
- ndim
- )));
- }
+ _ => unreachable!("validate_shape() should have caught invalid
ndim"),
}
}
- // Fallback: try to extract as Vec<f64> (Python list)
- if let Ok(vec_data) = data.extract::<Vec<f64>>() {
- let ptr = self
- .engine
- .encode(&vec_data, num_qubits, encoding_method)
- .map_err(|e| PyRuntimeError::new_err(format!("Encoding failed:
{}", e)))?;
- return Ok(QuantumTensor {
- ptr,
- consumed: false,
- });
+ // CPU tensor path
+ validate_tensor(data)?;
+ // PERF: Avoid Tensor -> Python list -> Vec deep copies.
+ //
+ // For CPU tensors, `tensor.detach().numpy()` returns a NumPy view
that shares the same
+ // underlying memory (zero-copy) when the tensor is C-contiguous. We
can then borrow a
+ // `&[f64]` directly via pyo3-numpy.
+ let ndim: usize = data.call_method0("dim")?.extract()?;
+ validate_shape(ndim, "tensor")?;
+ let numpy_view = data
+ .call_method0("detach")?
+ .call_method0("numpy")
+ .map_err(|_| {
+ PyRuntimeError::new_err(
+ "Failed to convert torch.Tensor to NumPy view. Ensure the
tensor is on CPU \
+ and does not require grad (try: tensor =
tensor.detach().cpu())",
+ )
+ })?;
+
+ match ndim {
+ 1 => {
+ // 1D tensor: single sample encoding
+ let array_1d =
numpy_view.extract::<PyReadonlyArray1<f64>>().map_err(|_| {
+ PyRuntimeError::new_err(
+ "Failed to extract NumPy view as float64 array. Ensure
dtype is float64 \
+ (try: tensor = tensor.to(torch.float64))",
+ )
+ })?;
+ let data_slice = array_1d.as_slice().map_err(|_| {
+ PyRuntimeError::new_err(
+ "Tensor must be contiguous (C-order) to get zero-copy
slice \
+ (try: tensor = tensor.contiguous())",
+ )
+ })?;
+ let ptr = self
+ .engine
+ .encode(data_slice, num_qubits, encoding_method)
+ .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ Ok(QuantumTensor {
+ ptr,
+ consumed: false,
+ })
+ }
+ 2 => {
+ // 2D tensor: batch encoding
+ let array_2d =
numpy_view.extract::<PyReadonlyArray2<f64>>().map_err(|_| {
+ PyRuntimeError::new_err(
+ "Failed to extract NumPy view as float64 array. Ensure
dtype is float64 \
+ (try: tensor = tensor.to(torch.float64))",
+ )
+ })?;
+ let shape = array_2d.shape();
+ let num_samples = shape[0];
+ let sample_size = shape[1];
+ let data_slice = array_2d.as_slice().map_err(|_| {
+ PyRuntimeError::new_err(
+ "Tensor must be contiguous (C-order) to get zero-copy
slice \
+ (try: tensor = tensor.contiguous())",
+ )
+ })?;
+ let ptr = self
+ .engine
+ .encode_batch(
+ data_slice,
+ num_samples,
+ sample_size,
+ num_qubits,
+ encoding_method,
+ )
+ .map_err(|e| PyRuntimeError::new_err(format!("Encoding
failed: {}", e)))?;
+ Ok(QuantumTensor {
+ ptr,
+ consumed: false,
+ })
+ }
+ _ => unreachable!("validate_shape() should have caught invalid
ndim"),
}
+ }
- Err(PyRuntimeError::new_err(
- "Unsupported data type. Expected: list, NumPy array, PyTorch
tensor, or file path",
- ))
+ /// Encode from Python list
+ fn encode_from_list(
+ &self,
+ data: &Bound<'_, PyAny>,
+ num_qubits: usize,
+ encoding_method: &str,
+ ) -> PyResult<QuantumTensor> {
+ let vec_data = data.extract::<Vec<f64>>().map_err(|_| {
+ PyRuntimeError::new_err(
+ "Unsupported data type. Expected: list, NumPy array, PyTorch
tensor, or file path",
+ )
+ })?;
+ let ptr = self
+ .engine
+ .encode(&vec_data, num_qubits, encoding_method)
+ .map_err(|e| PyRuntimeError::new_err(format!("Encoding failed:
{}", e)))?;
+ Ok(QuantumTensor {
+ ptr,
+ consumed: false,
+ })
}
/// Internal helper to encode from file based on extension
diff --git a/testing/qdp/test_bindings.py b/testing/qdp/test_bindings.py
index 3823db3d1..56b2da7b2 100644
--- a/testing/qdp/test_bindings.py
+++ b/testing/qdp/test_bindings.py
@@ -17,6 +17,7 @@
"""Simple tests for PyO3 bindings."""
import pytest
+import torch
import _qdp
@@ -138,43 +139,43 @@ def test_pytorch_integration():
@pytest.mark.gpu
-def test_pytorch_precision_float64():
- """Verify optional float64 precision produces complex128 tensors."""
[email protected](
+ "precision,expected_dtype",
+ [
+ ("float32", "complex64"),
+ ("float64", "complex128"),
+ ],
+)
+def test_precision(precision, expected_dtype):
+ """Test different precision settings produce correct output dtypes."""
pytest.importorskip("torch")
import torch
from _qdp import QdpEngine
- engine = QdpEngine(0, precision="float64")
+ engine = QdpEngine(0, precision=precision)
data = [1.0, 2.0, 3.0, 4.0]
qtensor = engine.encode(data, 2, "amplitude")
torch_tensor = torch.from_dlpack(qtensor)
- assert torch_tensor.dtype == torch.complex128
-
-
[email protected]
-def test_encode_tensor_cpu():
- """Test encoding from CPU PyTorch tensor (1D, single sample)."""
- pytest.importorskip("torch")
- import torch
- from _qdp import QdpEngine
-
- if not torch.cuda.is_available():
- pytest.skip("GPU required for QdpEngine")
-
- engine = QdpEngine(0)
- data = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)
- qtensor = engine.encode(data, 2, "amplitude")
-
- # Verify result
- torch_tensor = torch.from_dlpack(qtensor)
- assert torch_tensor.is_cuda
- assert torch_tensor.shape == (1, 4)
+ expected = getattr(torch, expected_dtype)
+ assert torch_tensor.dtype == expected, (
+ f"Expected {expected_dtype}, got {torch_tensor.dtype}"
+ )
@pytest.mark.gpu
-def test_encode_tensor_batch():
- """Test encoding from CPU PyTorch tensor (2D, batch encoding with
zero-copy)."""
[email protected](
+ "data_shape,expected_shape",
+ [
+ ([1.0, 2.0, 3.0, 4.0], (1, 4)), # 1D tensor -> single sample
+ (
+ [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0,
12.0]],
+ (3, 4),
+ ), # 2D tensor -> batch
+ ],
+)
+def test_encode_tensor_cpu(data_shape, expected_shape):
+ """Test encoding from CPU PyTorch tensor (1D or 2D, zero-copy)."""
pytest.importorskip("torch")
import torch
from _qdp import QdpEngine
@@ -183,19 +184,16 @@ def test_encode_tensor_batch():
pytest.skip("GPU required for QdpEngine")
engine = QdpEngine(0)
- # Create 2D tensor (batch_size=3, features=4)
- data = torch.tensor(
- [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]],
- dtype=torch.float64,
- )
- assert data.is_contiguous(), "Test tensor should be contiguous for
zero-copy"
+ data = torch.tensor(data_shape, dtype=torch.float64)
+ if len(data_shape) > 1:
+ assert data.is_contiguous(), "Test tensor should be contiguous for
zero-copy"
qtensor = engine.encode(data, 2, "amplitude")
# Verify result
torch_tensor = torch.from_dlpack(qtensor)
assert torch_tensor.is_cuda
- assert torch_tensor.shape == (3, 4), "Batch encoding should preserve batch
size"
+ assert torch_tensor.shape == expected_shape
@pytest.mark.gpu
@@ -255,34 +253,19 @@ def test_encode_errors():
@pytest.mark.gpu
-def test_encode_cuda_tensor_1d():
- """Test encoding from 1D CUDA tensor (single sample, zero-copy)."""
- pytest.importorskip("torch")
- import torch
- from _qdp import QdpEngine
-
- if not torch.cuda.is_available():
- pytest.skip("GPU required for QdpEngine")
-
- engine = QdpEngine(0)
-
- # Create 1D CUDA tensor
- data = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64,
device="cuda:0")
- qtensor = engine.encode(data, 2, "amplitude")
-
- # Verify result
- result = torch.from_dlpack(qtensor)
- assert result.is_cuda
- assert result.shape == (1, 4) # 2^2 = 4 amplitudes
-
- # Verify normalization (amplitudes should have unit norm)
- norm = torch.sqrt(torch.sum(torch.abs(result) ** 2))
- assert torch.isclose(norm, torch.tensor(1.0, device="cuda:0"), atol=1e-6)
-
-
[email protected]
-def test_encode_cuda_tensor_2d_batch():
- """Test encoding from 2D CUDA tensor (batch, zero-copy)."""
[email protected](
+ "data_shape,expected_shape,expected_batch_size",
+ [
+ ([1.0, 2.0, 3.0, 4.0], (1, 4), 1), # 1D tensor -> single sample
+ (
+ [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0,
12.0]],
+ (3, 4),
+ 3,
+ ), # 2D tensor -> batch
+ ],
+)
+def test_encode_cuda_tensor(data_shape, expected_shape, expected_batch_size):
+ """Test encoding from CUDA tensor (1D or 2D, zero-copy)."""
pytest.importorskip("torch")
import torch
from _qdp import QdpEngine
@@ -292,21 +275,17 @@ def test_encode_cuda_tensor_2d_batch():
engine = QdpEngine(0)
- # Create 2D CUDA tensor (batch_size=3, features=4)
- data = torch.tensor(
- [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]],
- dtype=torch.float64,
- device="cuda:0",
- )
+ # Create CUDA tensor
+ data = torch.tensor(data_shape, dtype=torch.float64, device="cuda:0")
qtensor = engine.encode(data, 2, "amplitude")
# Verify result
result = torch.from_dlpack(qtensor)
assert result.is_cuda
- assert result.shape == (3, 4) # batch_size=3, 2^2=4
+ assert result.shape == expected_shape
- # Verify each sample is normalized
- for i in range(3):
+ # Verify normalization (each sample should have unit norm)
+ for i in range(expected_batch_size):
norm = torch.sqrt(torch.sum(torch.abs(result[i]) ** 2))
assert torch.isclose(norm, torch.tensor(1.0, device="cuda:0"),
atol=1e-6)
@@ -389,8 +368,15 @@ def test_encode_cuda_tensor_empty():
@pytest.mark.gpu
-def test_encode_cuda_tensor_preserves_input():
- """Test that input CUDA tensor is not modified after encoding."""
[email protected](
+ "data_shape,is_batch",
+ [
+ ([1.0, 2.0, 3.0, 4.0], False), # 1D tensor
+ ([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], True), # 2D tensor
(batch)
+ ],
+)
+def test_encode_cuda_tensor_preserves_input(data_shape, is_batch):
+ """Test that input CUDA tensor (1D or 2D) is not modified after
encoding."""
pytest.importorskip("torch")
import torch
from _qdp import QdpEngine
@@ -401,8 +387,7 @@ def test_encode_cuda_tensor_preserves_input():
engine = QdpEngine(0)
# Create CUDA tensor and save a copy
- original_data = [1.0, 2.0, 3.0, 4.0]
- data = torch.tensor(original_data, dtype=torch.float64, device="cuda:0")
+ data = torch.tensor(data_shape, dtype=torch.float64, device="cuda:0")
data_clone = data.clone()
# Encode
@@ -414,7 +399,8 @@ def test_encode_cuda_tensor_preserves_input():
@pytest.mark.gpu
-def test_encode_cuda_tensor_unsupported_encoding():
[email protected]("encoding_method", ["basis", "angle"])
+def test_encode_cuda_tensor_unsupported_encoding(encoding_method):
"""Test error when using CUDA tensor with unsupported encoding method."""
pytest.importorskip("torch")
import torch
@@ -430,53 +416,23 @@ def test_encode_cuda_tensor_unsupported_encoding():
data = torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=torch.float64,
device="cuda:0")
with pytest.raises(RuntimeError, match="only supports 'amplitude' method"):
- engine.encode(data, 2, "basis")
-
- with pytest.raises(RuntimeError, match="only supports 'amplitude' method"):
- engine.encode(data, 2, "angle")
-
-
[email protected]
-def test_encode_cuda_tensor_3d_rejected():
- """Test error when CUDA tensor has 3+ dimensions."""
- pytest.importorskip("torch")
- import torch
- from _qdp import QdpEngine
-
- if not torch.cuda.is_available():
- pytest.skip("GPU required for QdpEngine")
-
- engine = QdpEngine(0)
-
- # Create 3D CUDA tensor (should be rejected)
- data = torch.randn(2, 3, 4, dtype=torch.float64, device="cuda:0")
- with pytest.raises(RuntimeError, match="Unsupported CUDA tensor shape:
3D"):
- engine.encode(data, 2, "amplitude")
-
-
[email protected]
-def test_encode_cuda_tensor_zero_values():
- """Test error when CUDA tensor contains all zeros (zero norm)."""
- pytest.importorskip("torch")
- import torch
- from _qdp import QdpEngine
-
- if not torch.cuda.is_available():
- pytest.skip("GPU required for QdpEngine")
-
- engine = QdpEngine(0)
-
- # Create CUDA tensor with all zeros (cannot be normalized)
- data = torch.zeros(4, dtype=torch.float64, device="cuda:0")
- with pytest.raises(RuntimeError, match="zero or non-finite norm"):
- engine.encode(data, 2, "amplitude")
+ engine.encode(data, 2, encoding_method)
@pytest.mark.gpu
-def test_encode_cuda_tensor_nan_values():
- """Test error when CUDA tensor contains NaN values."""
[email protected](
+ "input_type,error_match",
+ [
+ ("cuda_tensor", "Unsupported CUDA tensor shape: 3D"),
+ ("cpu_tensor", "Unsupported tensor shape: 3D"),
+ ("numpy_array", "Unsupported array shape: 3D"),
+ ],
+)
+def test_encode_3d_rejected(input_type, error_match):
+ """Test error when input has 3+ dimensions (CUDA tensor, CPU tensor, or
NumPy array)."""
pytest.importorskip("torch")
import torch
+ import numpy as np
from _qdp import QdpEngine
if not torch.cuda.is_available():
@@ -484,17 +440,41 @@ def test_encode_cuda_tensor_nan_values():
engine = QdpEngine(0)
- # Create CUDA tensor with NaN
- data = torch.tensor(
- [1.0, float("nan"), 3.0, 4.0], dtype=torch.float64, device="cuda:0"
- )
- with pytest.raises(RuntimeError, match="zero or non-finite norm"):
+ # Create 3D data based on input type
+ if input_type == "cuda_tensor":
+ data = torch.randn(2, 3, 4, dtype=torch.float64, device="cuda:0")
+ elif input_type == "cpu_tensor":
+ data = torch.randn(2, 3, 4, dtype=torch.float64)
+ elif input_type == "numpy_array":
+ data = np.random.randn(2, 3, 4).astype(np.float64)
+ else:
+ raise ValueError(f"Unknown input_type: {input_type}")
+
+ with pytest.raises(RuntimeError, match=error_match):
engine.encode(data, 2, "amplitude")
@pytest.mark.gpu
-def test_encode_cuda_tensor_inf_values():
- """Test error when CUDA tensor contains Inf values."""
[email protected](
+ "tensor_factory,description",
+ [
+ (lambda: torch.zeros(4, dtype=torch.float64, device="cuda:0"),
"zeros"),
+ (
+ lambda: torch.tensor(
+ [1.0, float("nan"), 3.0, 4.0], dtype=torch.float64,
device="cuda:0"
+ ),
+ "NaN",
+ ),
+ (
+ lambda: torch.tensor(
+ [1.0, float("inf"), 3.0, 4.0], dtype=torch.float64,
device="cuda:0"
+ ),
+ "Inf",
+ ),
+ ],
+)
+def test_encode_cuda_tensor_non_finite_values(tensor_factory, description):
+ """Test error when CUDA tensor contains non-finite values (zeros, NaN,
Inf)."""
pytest.importorskip("torch")
import torch
from _qdp import QdpEngine
@@ -503,17 +483,21 @@ def test_encode_cuda_tensor_inf_values():
pytest.skip("GPU required for QdpEngine")
engine = QdpEngine(0)
+ data = tensor_factory()
- # Create CUDA tensor with Inf
- data = torch.tensor(
- [1.0, float("inf"), 3.0, 4.0], dtype=torch.float64, device="cuda:0"
- )
with pytest.raises(RuntimeError, match="zero or non-finite norm"):
engine.encode(data, 2, "amplitude")
@pytest.mark.gpu
-def test_encode_cuda_tensor_output_dtype():
[email protected](
+ "precision,expected_dtype",
+ [
+ ("float32", torch.complex64),
+ ("float64", torch.complex128),
+ ],
+)
+def test_encode_cuda_tensor_output_dtype(precision, expected_dtype):
"""Test that CUDA tensor encoding produces correct output dtype."""
pytest.importorskip("torch")
import torch
@@ -522,45 +506,12 @@ def test_encode_cuda_tensor_output_dtype():
if not torch.cuda.is_available():
pytest.skip("GPU required for QdpEngine")
- # Test default precision (float32 -> complex64)
- engine_f32 = QdpEngine(0, precision="float32")
- data = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64,
device="cuda:0")
- result = torch.from_dlpack(engine_f32.encode(data, 2, "amplitude"))
- assert result.dtype == torch.complex64, f"Expected complex64, got
{result.dtype}"
-
- # Test float64 precision (float64 -> complex128)
- engine_f64 = QdpEngine(0, precision="float64")
+ engine = QdpEngine(0, precision=precision)
data = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64,
device="cuda:0")
- result = torch.from_dlpack(engine_f64.encode(data, 2, "amplitude"))
- assert result.dtype == torch.complex128, f"Expected complex128, got
{result.dtype}"
-
-
[email protected]
-def test_encode_cuda_tensor_preserves_input_batch():
- """Test that input 2D CUDA tensor (batch) is not modified after
encoding."""
- pytest.importorskip("torch")
- import torch
- from _qdp import QdpEngine
-
- if not torch.cuda.is_available():
- pytest.skip("GPU required for QdpEngine")
-
- engine = QdpEngine(0)
-
- # Create 2D CUDA tensor and save a copy
- data = torch.tensor(
- [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
- dtype=torch.float64,
- device="cuda:0",
+ result = torch.from_dlpack(engine.encode(data, 2, "amplitude"))
+ assert result.dtype == expected_dtype, (
+ f"Expected {expected_dtype}, got {result.dtype}"
)
- data_clone = data.clone()
-
- # Encode
- qtensor = engine.encode(data, 2, "amplitude")
- _ = torch.from_dlpack(qtensor)
-
- # Verify original tensor is unchanged
- assert torch.equal(data, data_clone)
@pytest.mark.gpu
@@ -766,7 +717,70 @@ def test_angle_encode_errors():
engine.encode([float("nan"), 0.0], 2, "angle")
-# ==================== IQP Encoding Tests ====================
[email protected]
[email protected](
+ "data_shape,expected_shape",
+ [
+ ([1.0, 2.0, 3.0, 4.0], (1, 4)), # 1D array -> single sample
+ (
+ [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
+ (2, 4),
+ ), # 2D array -> batch
+ ],
+)
+def test_encode_numpy_array(data_shape, expected_shape):
+ """Test encoding from NumPy array (1D or 2D)."""
+ pytest.importorskip("torch")
+ import numpy as np
+ import torch
+ from _qdp import QdpEngine
+
+ if not torch.cuda.is_available():
+ pytest.skip("GPU required for QdpEngine")
+
+ engine = QdpEngine(0)
+ data = np.array(data_shape, dtype=np.float64)
+ qtensor = engine.encode(data, 2, "amplitude")
+
+ # Verify result
+ torch_tensor = torch.from_dlpack(qtensor)
+ assert torch_tensor.is_cuda
+ assert torch_tensor.shape == expected_shape
+
+
[email protected]
+def test_encode_pathlib_path():
+ """Test encoding from pathlib.Path object."""
+ pytest.importorskip("torch")
+ import numpy as np
+ import torch
+ from pathlib import Path
+ import tempfile
+ import os
+ from _qdp import QdpEngine
+
+ if not torch.cuda.is_available():
+ pytest.skip("GPU required for QdpEngine")
+
+ engine = QdpEngine(0)
+ num_qubits = 2
+ sample_size = 2**num_qubits
+
+ # Create temporary .npy file
+ data = np.array([[1.0, 2.0, 3.0, 4.0], [0.5, 0.5, 0.5, 0.5]],
dtype=np.float64)
+ with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
+ npy_path = Path(f.name)
+ np.save(npy_path, data)
+
+ try:
+ # Test with pathlib.Path
+ qtensor = engine.encode(npy_path, num_qubits, "amplitude")
+ torch_tensor = torch.from_dlpack(qtensor)
+ assert torch_tensor.is_cuda
+ assert torch_tensor.shape == (2, sample_size)
+ finally:
+ if os.path.exists(npy_path):
+ os.remove(npy_path)
@pytest.mark.gpu
