torchrl-nightly 2025.8.12__cp39-cp39-win_amd64.whl → 2025.8.13__cp39-cp39-win_amd64.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.12.dist-info → torchrl_nightly-2025.8.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchrl-nightly
-Version: 2025.8.12
+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.12.dist-info → torchrl_nightly-2025.8.13.dist-info}/RECORD

@@ -92,7 +92,7 @@ sota-implementations/td3_bc/td3_bc.py,sha256=GUGzbGkaR7MDbmhvMCs69Ewd-iCVjaQCEt6
 sota-implementations/td3_bc/utils.py,sha256=ZMMYehJP6KZaiiwTRlJRFKC528M49QhfEb_RAxaBj-U,6951
 torchrl/__init__.py,sha256=NjAWriDfFsQIxAUG6lvmqJnZE2olvHF2LDN9ZE1WT9M,3065
 torchrl/_extension.py,sha256=5DqUOUHZJPLqcSvGztUUBy4rJFAzznuJCAu5c84OJo8,2981
-torchrl/_torchrl.cp39-win_amd64.pyd,sha256=qhYN9yg8Nf-78WrcGTgXn0zw4kChL8GdHujCWBOIZM0,452608
+torchrl/_torchrl.cp39-win_amd64.pyd,sha256=D1Hf5b_m4RGxiUAJlNXDdc0MvAWLjNtLCfiesTL6pu4,452608
 torchrl/_utils.py,sha256=9s4EpAZlaUVhWb7Llvg86Pex0sRv48DDQia30VbvKvo,42339
 torchrl/collectors/__init__.py,sha256=LzTyfxmkNGPSa5-3rS5unQK7HfT5ZEdr2NV291rAOlU,832
 torchrl/collectors/collectors.py,sha256=XsfSHhSPWi_Czt-Olpgr6hgukXJM-1oVuWJ1m4yfZX4,182174
@@ -180,7 +180,7 @@ torchrl/envs/libs/gym.py,sha256=4hWb06j5DFSU7S6k9wsXmuwtCsUVK3WtGYqsHregEcg,8255
 torchrl/envs/libs/habitat.py,sha256=w_B1VQQRcjH0kbynu6H9zcimjoluEqVcsH1P0WuQguM,5525
 torchrl/envs/libs/isaac_lab.py,sha256=lD3PVhuco2mncgSK-714yPQ9EjxdLOKRfDRR__eARec,3523
 torchrl/envs/libs/isaacgym.py,sha256=C9fY_gLTRm1qk0TGMFMQKXNS27HBpIqRz2rjdJ0Lg7k,7336
-torchrl/envs/libs/jax_utils.py,sha256=j63ti4q7owJf73C-h5Wi9AoglVlQrTyABJ9fCV9RfbM,6201
+torchrl/envs/libs/jax_utils.py,sha256=q06IOxSX0jCfm1JnVGj-SldYHtUrlVoKtZj8U1zDJcI,6247
 torchrl/envs/libs/jumanji.py,sha256=6C9NmILOt3wbgoEB4oHaDU-4lBRzJGQKJ62hYU8GLTg,41234
 torchrl/envs/libs/meltingpot.py,sha256=5IT7Kj5C6u1CWDAdb1SimCgEv33ti1Y7hVsJy5Nf8bA,26695
 torchrl/envs/libs/openml.py,sha256=DDG5qbBe7ipjQ8nJUg6pBVja8w9zsRJFY6n_m10zItg,5866
@@ -322,8 +322,8 @@ torchrl/trainers/helpers/losses.py,sha256=HwrovwbMOhY-5-hlOz-YHclKnoJhMijVjDNuAT
 torchrl/trainers/helpers/models.py,sha256=VujBq9H92sEzpCtU1iTrJQNlwvyOO-Rho4bzsMonX6s,22465
 torchrl/trainers/helpers/replay_buffer.py,sha256=RaZqXnHimmadiibvDBcLbtIhpPaVMTPhYMOBvX4v3CA,2060
 torchrl/trainers/helpers/trainers.py,sha256=hB1FtHtP-S0PBQ4LF6WPy37caaLpacyaLThj1BNl5Ho,12372
-torchrl_nightly-2025.8.12.dist-info/LICENSE,sha256=PGO-oZsq4EzhE1-WQS2xGiEF3UCVb9YawfQ09cIMV_8,1119
-torchrl_nightly-2025.8.12.dist-info/METADATA,sha256=k896JgYWNPyT4FIriOiAYTlQ-56mzdVKRO8zGJhYhYw,42540
-torchrl_nightly-2025.8.12.dist-info/WHEEL,sha256=yA7mxgqX2UV73NtJdMh2AAmdb628loM81912H3s5r00,100
-torchrl_nightly-2025.8.12.dist-info/top_level.txt,sha256=-5FcSdmJ9DwdHF8aOIaofsPbz4Gm8G1eo7r7Sc2CHgE,59
-torchrl_nightly-2025.8.12.dist-info/RECORD,,
+torchrl_nightly-2025.8.13.dist-info/LICENSE,sha256=PGO-oZsq4EzhE1-WQS2xGiEF3UCVb9YawfQ09cIMV_8,1119
+torchrl_nightly-2025.8.13.dist-info/METADATA,sha256=WPKgQIOQ9ox9gxi54abdpV03Bi-d9ppl58i9vKZVqmk,42652
+torchrl_nightly-2025.8.13.dist-info/WHEEL,sha256=yA7mxgqX2UV73NtJdMh2AAmdb628loM81912H3s5r00,100
+torchrl_nightly-2025.8.13.dist-info/top_level.txt,sha256=-5FcSdmJ9DwdHF8aOIaofsPbz4Gm8G1eo7r7Sc2CHgE,59
+torchrl_nightly-2025.8.13.dist-info/RECORD,,