fix: update JAX DLPack API to remove deprecation warning

[adamimos]

The jax.dlpack.to_dlpack function was deprecated in JAX v0.6.0 and removed in v0.7.0.

Updated jax_to_torch function to pass JAX arrays directly to torch.utils.dlpack.from_dlpack, which uses the newer dlpack protocol automatically.

Fixes #70

Generated with Claude Code

[claude]

Code Review

Overall Assessment: LGTM with minor observations

This is a solid fix that addresses the JAX DLPack API deprecation warning properly.

Code Quality & Best Practices

• Excellent: The solution follows the newer DLPack protocol correctly by passing JAX arrays directly to torch.utils.dlpack.from_dlpack
• Good: Maintains the existing fallback mechanism to numpy conversion
• Good: Preserves all error handling and warning patterns
• Good: Correctly updated the docstring to remove the obsolete device parameter

Technical Correctness

• API Migration: Correctly removes the deprecated jax.dlpack.to_dlpack() call
• Protocol Compliance: Uses the newer dlpack protocol automatically through direct array passing
• Backward Compatibility: Maintains the same function signature and behavior
• Error Handling: Preserves existing try/catch with numpy fallback

Performance Considerations

• Improvement: Eliminates an intermediate DLPack tensor creation, potentially slightly more efficient
• Maintained: Zero-copy GPU transfers still work as expected
• No Regression: Fallback path unchanged for cases where DLPack fails

Security & Safety

• No Security Concerns: Change is purely API migration
• Type Safety: Maintains existing type annotations
• Exception Handling: Preserved comprehensive error handling

Test Coverage

• Good: Existing tests in tests/utils/test_pytorch_utils.py should cover this change
• Recommendation: Consider running the test suite to verify GPU/CPU conversion still works correctly

Minor Observations

1. Unused Import: The 'from jax import dlpack as jax_dlpack' import on line 13 is now unused and could be removed
2. Documentation: The module docstring mentions DLPack but could be updated to reflect the newer protocol usage

Code Style Compliance

• Line Length: Within 120 character limit
• Type Hints: Properly maintained
• Functional Style: Maintains pure function approach
• Naming: Follows project conventions

Recommendation

This PR should be merged after addressing the unused import. The fix correctly resolves the deprecation warning while maintaining full functionality.

[adamimos]

Code Review

Overall Assessment: LGTM with minor observations

This is a solid fix that addresses the JAX DLPack API deprecation warning properly.

Code Quality & Best Practices

• Excellent: The solution follows the newer DLPack protocol correctly by passing JAX arrays directly to torch.utils.dlpack.from_dlpack
• Good: Maintains the existing fallback mechanism to numpy conversion
• Good: Preserves all error handling and warning patterns
• Good: Correctly updated the docstring to remove the obsolete device parameter

Technical Correctness

• API Migration: Correctly removes the deprecated jax.dlpack.to_dlpack() call
• Protocol Compliance: Uses the newer dlpack protocol automatically through direct array passing
• Backward Compatibility: Maintains the same function signature and behavior
• Error Handling: Preserves existing try/catch with numpy fallback

Performance Considerations

• Improvement: Eliminates an intermediate DLPack tensor creation, potentially slightly more efficient
• Maintained: Zero-copy GPU transfers still work as expected
• No Regression: Fallback path unchanged for cases where DLPack fails

Security & Safety

• No Security Concerns: Change is purely API migration
• Type Safety: Maintains existing type annotations
• Exception Handling: Preserved comprehensive error handling

Test Coverage

• Good: Existing tests in tests/utils/test_pytorch_utils.py should cover this change
• Recommendation: Consider running the test suite to verify GPU/CPU conversion still works correctly

Minor Observations

1. Unused Import: The 'from jax import dlpack as jax_dlpack' import on line 13 is now unused and could be removed
2. Documentation: The module docstring mentions DLPack but could be updated to reflect the newer protocol usage

Code Style Compliance

• Line Length: Within 120 character limit
• Type Hints: Properly maintained
• Functional Style: Maintains pure function approach
• Naming: Follows project conventions

Recommendation

This PR should be merged after addressing the unused import. The fix correctly resolves the deprecation warning while maintaining full functionality.

@claude please make changes to address the unused import

[claude]

Claude finished @adamimos's task —— View job

---

Addressing Unused Import

Todo List:

• Read the current pytorch_utils.py file to confirm the unused import
• Verify the import usage in the codebase
• Update comment with findings

Findings:

After reviewing the code more carefully, I need to correct my previous assessment. The from jax import dlpack as jax_dlpack import on line 13 is NOT actually unused.

While the PR correctly removed the usage of jax_dlpack.to_dlpack() from the jax_to_torch function (which was deprecated), the jax_dlpack import is still needed for the torch_to_jax function on line 69:

jax_array = jax_dlpack.from_dlpack(dlpack_tensor)

The jax_dlpack.from_dlpack() function is still the correct way to convert from DLPack tensors to JAX arrays, so this import should remain.

Conclusion:

No changes are needed. My previous code review incorrectly identified this as an unused import. The import is actively used and required for the torch_to_jax function to work properly.

[ealt]

Good job, Claude!