[QDP] Fix benchmark e2e Error

[400Ping]

Purpose of PR

DLPack output was coming back as a flat 1-D tensor (n_samples * 2^n_qubits), so the benchmark assert expected 2-D [n_samples, state_len] was failing. Added a defensive reshape to view the tensor as [n_samples, state_len] before asserting, so benchmarks can continue even when the producer flattens the DLPack shape.

Error log:

$ python ./benchmark/benchmark_e2e.py --frameworks all --samples 1000
Generating 1000 samples of 16 qubits...
  Generated 1000 samples
  Parquet: 500.34 MB, Arrow IPC: 500.00 MB

======================================================================
E2E BENCHMARK: 16 Qubits, 1000 Samples
======================================================================

[PennyLane] Full Pipeline (Disk -> GPU)...
  IO Time: 0.9895 s
/home/jay/work/mahout/qdp/./benchmark/benchmark_e2e.py:237: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at /pytorch/aten/src/ATen/native/Copy.cpp:309.)
  state_gpu = state_cpu.to("cuda", dtype=torch.float32)
  Total Time: 1.4611 s

[Qiskit] Full Pipeline (Disk -> GPU)...
  IO Time: 0.9268 s
    Processed 40/40 vectors...
  Total Time: 68.6920 s

[Mahout-Parquet] Full Pipeline (Parquet -> GPU)...
  Parquet->GPU (IO+Encode): 0.7504 s
  DLPack conversion: 0.0001 s
Traceback (most recent call last):
  File "/home/jay/work/mahout/qdp/./benchmark/benchmark_e2e.py", line 450, in <module>
    t_mahout_parquet, mahout_parquet_all_states = run_mahout_parquet(
                                                  ^^^^^^^^^^^^^^^^^^^
  File "/home/jay/work/mahout/qdp/./benchmark/benchmark_e2e.py", line 282, in run_mahout_parquet
    assert gpu_batched.shape == (n_samples, state_len), (

Related Issues or PRs

Closes #759

Changes Made

• Bug fix
• New feature
• Refactoring
• Documentation
• Test
• CI/CD pipeline
• Other

Breaking Changes

• Yes
• No

Checklist

• Added or updated unit tests for all changes
• Added or updated documentation for all changes
• Successfully built and ran all unit tests or manual tests locally
• PR title follows "MAHOUT-XXX: Brief Description" format (if related to an issue)
• Code follows ASF guidelines

[400Ping]

cc @rich7420 @guan404ming @ryankert01

[guan404ming]

My setup can work without this fix. I’ll leave it as is and let someone else take a look. Thanks.

[400Ping]

May I ask what your pytorch version is? I think it is because of this.

[400Ping]

nvm, fixed my env. Closing this pr