Fix JAX DLPack export to use __dlpack__ protocol directly

JanuszL · JanuszL · commit bb3fb172c5ec · 2026-03-09T07:51:25.000+01:00
- Replaces deprecated jax.dlpack.to_dlpack() calls with the standard
  tensor.__dlpack__() method, which is the correct DLPack protocol
  interface for JAX 0.6+.

Signed-off-by: Janusz Lisiecki &lt;jlisiecki@nvidia.com&gt;
diff --git a/dali/python/nvidia/dali/plugin/jax/fn/_function_transform.py b/dali/python/nvidia/dali/plugin/jax/fn/_function_transform.py
@@ -29,7 +29,7 @@ def gpu_to_dlpack(tensor: jax.Array, stream):
             f"The function returned array residing on the device of "
             f"kind `{devices[0].platform}`, expected `gpu`."
         )
-    return jax.dlpack.to_dlpack(tensor, stream=stream)
+    return tensor.__dlpack__(stream=stream)
 
 
 def cpu_to_dlpack(tensor: jax.Array):
@@ -44,7 +44,7 @@ def cpu_to_dlpack(tensor: jax.Array):
             f"The function returned array residing on the device of "
             f"kind `{devices[0].platform}`, expected `cpu`."
         )
-    return jax.dlpack.to_dlpack(tensor)
+    return tensor.__dlpack__()
 
 
 def with_gpu_dl_tensors_as_arrays(callback):

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ def gpu_to_dlpack(tensor: jax.Array, stream):`
`29`	`29`	`f"The function returned array residing on the device of "`
`30`	`30`	f"kind `{devices[0].platform}`, expected `gpu`."
`31`	`31`	`)`
`32`		`- return jax.dlpack.to_dlpack(tensor, stream=stream)`
	`32`	`+ return tensor.__dlpack__(stream=stream)`
`33`	`33`
`34`	`34`
`35`	`35`	`def cpu_to_dlpack(tensor: jax.Array):`
`@@ -44,7 +44,7 @@ def cpu_to_dlpack(tensor: jax.Array):`
`44`	`44`	`f"The function returned array residing on the device of "`
`45`	`45`	f"kind `{devices[0].platform}`, expected `cpu`."
`46`	`46`	`)`
`47`		`- return jax.dlpack.to_dlpack(tensor)`
	`47`	`+ return tensor.__dlpack__()`
`48`	`48`
`49`	`49`
`50`	`50`	`def with_gpu_dl_tensors_as_arrays(callback):`