torchrl-nightly 2025.8.11__cp39-cp39-macosx_10_9_universal2.whl → 2025.8.13__cp39-cp39-macosx_10_9_universal2.whl

torchrl/envs/libs/jax_utils.py

@@ -42,13 +42,13 @@ _dtype_conversion = {
 
 
 def _ndarray_to_tensor(value: jnp.ndarray | np.ndarray) -> torch.Tensor:  # noqa: F821
-    from jax import dlpack as jax_dlpack, numpy as jnp
+    from jax import numpy as jnp
 
     # JAX arrays generated by jax.vmap would have Numpy dtypes.
     if value.dtype in _dtype_conversion:
         value = value.view(_dtype_conversion[value.dtype])
     if isinstance(value, jnp.ndarray):
-        dlpack_tensor = jax_dlpack.to_dlpack(value)
+        dlpack_tensor = value.__dlpack__()
     elif isinstance(value, np.ndarray):
         dlpack_tensor = value.__dlpack__()
     else:
@@ -61,7 +61,9 @@ def _ndarray_to_tensor(value: jnp.ndarray | np.ndarray) -> torch.Tensor:  # noqa
 def _tensor_to_ndarray(value: torch.Tensor) -> jnp.ndarray:  # noqa: F821
     from jax import dlpack as jax_dlpack
 
-    return jax_dlpack.from_dlpack(torch_dlpack.to_dlpack(value.contiguous()))
+    # Detach the tensor to remove gradients before converting to DLPack
+    value = value.contiguous().detach()
+    return jax_dlpack.from_dlpack(value)
 
 
 def _get_object_fields(obj) -> dict:

{torchrl_nightly-2025.8.11.dist-info → torchrl_nightly-2025.8.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchrl-nightly
-Version: 2025.8.11
+Version: 2025.8.13
 Summary: A modular, primitive-first, python-first PyTorch library for Reinforcement Learning
 Author-email: torchrl contributors <vmoens@fb.com>
 Maintainer-email: torchrl contributors <vmoens@fb.com>
@@ -33,6 +33,9 @@ Requires-Dist: cloudpickle
 Requires-Dist: tensordict-nightly
 Provides-Extra: atari
 Requires-Dist: gymnasium[atari]; extra == "atari"
+Provides-Extra: brax
+Requires-Dist: jax[cuda12]>=0.7.0; extra == "brax"
+Requires-Dist: brax; extra == "brax"
 Provides-Extra: checkpointing
 Requires-Dist: torchsnapshot; extra == "checkpointing"
 Provides-Extra: dev

{torchrl_nightly-2025.8.11.dist-info → torchrl_nightly-2025.8.13.dist-info}/RECORD

@@ -92,7 +92,7 @@ sota-implementations/td3_bc/td3_bc.py,sha256=-e900G5UFkUnU9k9LNvJ1Qi-rHIk5o4WkhZ
 sota-implementations/td3_bc/utils.py,sha256=a4ZCpaIvXDra1UGC8eWS-T099-8OAVr6jsyDvfUN_IY,6700
 torchrl/__init__.py,sha256=0gcnqA2NCgcRYgg12F0lYJknQkbOwLRx__sgCahqARI,2949
 torchrl/_extension.py,sha256=0A19lzoHIZ4SNtoRlF6pwV1_k4qNocE2SsBhD7YlMJE,2917
-torchrl/_torchrl.cpython-39-darwin.so,sha256=t8Nau2UctljZORDsSXmqYruXU5eXlgQ1hqQr0zBAsjY,1632656
+torchrl/_torchrl.cpython-39-darwin.so,sha256=rzmPhvGyS2PtuSZ_CiI62qoRmCnvNC8g3BR3AZGtwzo,1632656
 torchrl/_utils.py,sha256=93LEupTGxX7-8U-zj291f1TEoGLdkiaHOeTxLxixKsY,41119
 torchrl/collectors/__init__.py,sha256=hJ3JD6shRku0BL6SzJQq44FZ5Q1RGR8LealFyU3FRn4,799
 torchrl/collectors/collectors.py,sha256=dTRCDfeDzZCkGNTbX1zrCfFGWp70-DTUjIaNYMJh9rQ,178314
@@ -180,7 +180,7 @@ torchrl/envs/libs/gym.py,sha256=b01bW-ostmOdcLBIl03Pcv4aWcJ2u4brYdyHezXZKDI,8045
 torchrl/envs/libs/habitat.py,sha256=PYvVqS8AGwplSY9r1x6RSLPejSKvxEyZgcDXFZvCYqA,5387
 torchrl/envs/libs/isaac_lab.py,sha256=4QiCB3tVGQ_D-gOngHyhtVrivy76B32fKGvpt-0hwxo,3436
 torchrl/envs/libs/isaacgym.py,sha256=GqXwbs9Iq7A5PHr1Sa-7UgQF42DLW7BFcu8lWYhBVJ8,7133
-torchrl/envs/libs/jax_utils.py,sha256=L_P9cd92zO08teZ-lKR2DF2hXk1_rzVIe1iBkLU7tsU,6037
+torchrl/envs/libs/jax_utils.py,sha256=MAm0dvsG6Uk73Obh5ablkGIgkdUgGZnOqqniST0-lZk,6081
 torchrl/envs/libs/jumanji.py,sha256=XS2dw3ekt8621VmMbASlUTyqHhLpzVgmuMyLgj07lks,40271
 torchrl/envs/libs/meltingpot.py,sha256=nd8P3JW_1D6fm_-eZbgmZle20fBoN6M_08SH5ydNS3Y,26096
 torchrl/envs/libs/openml.py,sha256=cuTWhedmXiDYjz8O-wsAuCKPfxaRGYaOffopGzidUw8,5713
@@ -322,8 +322,8 @@ torchrl/trainers/helpers/losses.py,sha256=sHlJqjh02t8cKN73X35Azd_OoWGurohLuviB8Y
 torchrl/trainers/helpers/models.py,sha256=ihTERG2c96E8cS3Tnul6a_ys6iDEEJmHh05p9blQTW8,21807
 torchrl/trainers/helpers/replay_buffer.py,sha256=ZUZHOa0TILyeWJ3iahzTJ6UvMl_0FdxuZfJEja94Bn8,2001
 torchrl/trainers/helpers/trainers.py,sha256=j6B5XA7_FFHMQeOIQwjNcO0CGE_4mZKUC9_jH_iqqh4,12071
-torchrl_nightly-2025.8.11.dist-info/LICENSE,sha256=xdjS4_xk-IwnLuIFCvTYTl9Y8aXRejqpmke3dGam_nI,1098
-torchrl_nightly-2025.8.11.dist-info/METADATA,sha256=cz_yY-HvDNgn25218PZ5-CTAZrlihMEOmOphpPTX6B4,41390
-torchrl_nightly-2025.8.11.dist-info/WHEEL,sha256=Kpi8Y7ETJ76qV_mhWFgtod7iK9qT4KSd2dlvMqnPXEY,113
-torchrl_nightly-2025.8.11.dist-info/top_level.txt,sha256=-5FcSdmJ9DwdHF8aOIaofsPbz4Gm8G1eo7r7Sc2CHgE,59
-torchrl_nightly-2025.8.11.dist-info/RECORD,,
+torchrl_nightly-2025.8.13.dist-info/LICENSE,sha256=xdjS4_xk-IwnLuIFCvTYTl9Y8aXRejqpmke3dGam_nI,1098
+torchrl_nightly-2025.8.13.dist-info/METADATA,sha256=0zoKgbG9IK7SgQb_1SWxF3PQUs5NV_GuxafkJwjUrF8,41499
+torchrl_nightly-2025.8.13.dist-info/WHEEL,sha256=Kpi8Y7ETJ76qV_mhWFgtod7iK9qT4KSd2dlvMqnPXEY,113
+torchrl_nightly-2025.8.13.dist-info/top_level.txt,sha256=-5FcSdmJ9DwdHF8aOIaofsPbz4Gm8G1eo7r7Sc2CHgE,59
+torchrl_nightly-2025.8.13.dist-info/RECORD,,