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
4 changes: 2 additions & 2 deletions qdp/qdp-python/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

use crate::pytorch::{
extract_cuda_tensor_info, get_torch_cuda_stream_ptr, is_cuda_tensor, is_pytorch_tensor,
validate_cuda_tensor_for_encoding, validate_shape, validate_tensor,
validate_cuda_tensor_for_encoding, validate_shape, validate_tensor_cpu,
};
use crate::tensor::QuantumTensor;
use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods};
Expand Down Expand Up @@ -206,7 +206,7 @@ impl QdpEngine {
}

// CPU tensor path
validate_tensor(data)?;
validate_tensor_cpu(data)?;
// PERF: Avoid Tensor -> Python list -> Vec deep copies.
//
// For CPU tensors, `tensor.detach().numpy()` returns a NumPy view that shares the same
Expand Down
12 changes: 6 additions & 6 deletions qdp/qdp-python/src/pytorch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,12 @@ pub fn is_pytorch_tensor(obj: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(module_name == "torch")
}

/// Helper to validate CPU tensor
pub fn validate_tensor(tensor: &Bound<'_, PyAny>) -> PyResult<()> {
if !is_pytorch_tensor(tensor)? {
return Err(PyRuntimeError::new_err("Object is not a PyTorch Tensor"));
}

/// Validate that a PyTorch tensor lives on the CPU.
///
/// The caller must have already confirmed `tensor` is a `torch.Tensor` (e.g.
/// via [`is_pytorch_tensor`]). This function only checks device placement, so
/// it also rejects non-CPU backends such as CUDA, MPS, XLA, and HPU.
pub fn validate_tensor_cpu(tensor: &Bound<'_, PyAny>) -> PyResult<()> {
let device = tensor.getattr("device")?;
let device_type: String = device.getattr("type")?.extract()?;

Expand Down
Loading