Skip to content

feat(qdp): add Jax support and generic DLPack fallbackc#1152

Closed
Sumit6307 wants to merge 1 commit intoapache:mainfrom
Sumit6307:feat/jax-dlpack-support
Closed

feat(qdp): add Jax support and generic DLPack fallbackc#1152
Sumit6307 wants to merge 1 commit intoapache:mainfrom
Sumit6307:feat/jax-dlpack-support

Conversation

@Sumit6307
Copy link

This PR aligns with the 2026 Q1 Roadmap for QDP (Quantum Data Plane) regarding "Integration & Release" and "Input Format Support."

Changes

  • New feature
  • Refactoring
  • Test

Why
Currently, the Quantum Data Plane (QDP) has tight integration with PyTorch but lacks native support for the Jax ecosystem, which is a leading framework for quantum machine learning research. Furthermore, the encode method was specialized for specific frameworks, making it harder to support other DLPack-compliant libraries like Cupy. This change enables Mahout to act as a truly universal Quantum Data Plane.

How

  • Jax Integration: Created a new jax.rs module that detects Jax arrays and handles asynchronous synchronization via block_until_ready() before data extraction.

  • Generic DLPack Fallback: Updated the unified encode method in engine.rs to support any Python object that implements the standard DLPack protocol (__dlpack__ and __dlpack_device__).

  • Framework Interoperability: Ensured that the logic is "non-invasive"—the QDP only attempts to synchronize Jax/DLPack objects when they are passed as input, maintaining zero overhead for other paths.

  • Improved Validation: Added robust error handling during the DLPack extraction process to ensure device and memory alignment consistency.

  • Tests: Added a new test suite in testing/qdp/test_jax.py specifically for Jax-to-QDP-to-PyTorch workflows.

@Sumit6307
Copy link
Author

@guan404ming Sir, Plz look into this PR

@400Ping 400Ping self-assigned this Mar 8, 2026
@400Ping
Copy link
Member

400Ping commented Mar 8, 2026

cc @viiccwen to take a look first before I make the final call

@400Ping
Copy link
Member

400Ping commented Mar 8, 2026

Have a better idea for jax, closing it.

@400Ping 400Ping closed this Mar 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants