API change for DLPack conversion

[adamimos]

I have been getting the following error:

packages/simplexity/generative_processes/torch_generator.py:32: UserWarning: DLPack conversion failed (jax.dlpack.to_dlpack was deprecated in JAX v0.6.0 and removed in JAX v0.7.0. Please use the newer DLPack API based on __dlpack__ and __dlpack_device__ instead. Typically, you can pass a JAX array directly to the `from_dlpack` function of another framework without using `to_dlpack`.), falling back to numpy. This may cause GPU-to-CPU transfer.
  return gen_states, jax_to_torch(inputs), jax_to_torch(labels)

This should be an easy API change

[adamimos]

@claude can you make a PR to fix this issue?

[claude]

Claude finished @adamimos's task —— View job • claude/issue-70-20250908-1901 • Create PR ➔

---

Fixing DLPack API Deprecation ✅

I've successfully fixed the DLPack API deprecation warning by updating the jax_to_torch function to use the newer JAX DLPack API.

Todo List:

• Read the current implementation in torch_generator.py
• Understand how jax_to_torch function works
• Update to use the newer DLPack API (__dlpack__ and __dlpack_device__)
• Commit and push changes
• Provide PR creation link

Changes Made:

• Updated simplexity/utils/pytorch_utils.py:38 to pass JAX arrays directly to torch_dlpack.from_dlpack(jax_array) instead of using the deprecated jax_dlpack.to_dlpack(jax_array)
• Removed the intermediate DLPack tensor creation
• Updated docstring to remove stale parameter documentation

Why This Fixes The Issue:
The new approach leverages JAX's newer DLPack protocol (__dlpack__ and __dlpack_device__) automatically when the JAX array is passed directly to PyTorch's from_dlpack function. This eliminates the deprecation warning while maintaining the same zero-copy GPU-to-GPU transfer functionality.

---

Job run | Branch