python312Packages.tensordict: 0.4.0 -> 0.5.0

Diff: https://github.com/pytorch/tensordict/compare/refs/tags/v0.4.0...v0.5.0

Changelog: https://github.com/pytorch/tensordict/releases/tag/v0.5.0

pkgs/development/python-modules/tensordict/default.nix

@@ -3,21 +3,27 @@
   buildPythonPackage,
   pythonOlder,
   fetchFromGitHub,
+
+  # build-system
   setuptools,
   torch,
-  wheel,
   which,
+
+  # dependencies
   cloudpickle,
   numpy,
+  orjson,
+
+  # checks
   h5py,
   pytestCheckHook,
+
   stdenv,
-  pythonAtLeast,
 }:
 
 buildPythonPackage rec {
   pname = "tensordict";
-  version = "0.4.0";
+  version = "0.5.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -26,19 +32,19 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "tensordict";
     rev = "refs/tags/v${version}";
-    hash = "sha256-wKEzNaaazGEkoElzp93RIlq/r5uRUdM7UyDy/DygIEc=";
+    hash = "sha256-jnRlN9gefR77pioIXf0qM1CP6EtpeQkBvVIecGkb/pw=";
   };
 
   build-system = [
     setuptools
     torch
-    wheel
     which
   ];
 
   dependencies = [
     cloudpickle
     numpy
+    orjson
     torch
   ];
 
@@ -55,8 +61,22 @@ buildPythonPackage rec {
   ];
 
   disabledTests =
-    # Hangs forever
-    [ "test_copy_onto" ]
+    [
+      # Hangs forever
+      "test_copy_onto"
+
+      # EOFError (MPI related)
+      # AssertionError: assert tensor(False)
+      # +  where tensor(False) = <built-in method all of Tensor object at 0x7ffe49bf87d0>()
+      "test_mp"
+
+      # torch._dynamo.exc.BackendCompilerFailed
+      # Requires a more recent version of triton
+      # Re-enable when https://github.com/NixOS/nixpkgs/pull/328247 is merged
+      "test_linear"
+      "test_seq"
+      "test_seq_lmbda"
+    ]
     ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-linux") [
       # RuntimeError: internal error
       "test_add_scale_sequence"
@@ -65,10 +85,19 @@ buildPythonPackage rec {
 
       # _queue.Empty errors in multiprocessing tests
       "test_isend"
+
+      # hangs forever
+      "test_map_iter_interrupt_early"
     ];
 
-  # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
-  disabledTestPaths = lib.optionals stdenv.isDarwin [ "test/test_distributed.py" ];
+  disabledTestPaths = lib.optionals stdenv.isDarwin [
+    # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
+    # OpenMP support not found.
+    "test/test_compile.py"
+
+    # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
+    "test/test_distributed.py"
+  ];
 
   meta = {
     description = "Pytorch dedicated tensor container";
@@ -76,7 +105,5 @@ buildPythonPackage rec {
     homepage = "https://github.com/pytorch/tensordict";
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [ GaetanLepage ];
-    # No python 3.12 support yet: https://github.com/pytorch/rl/issues/2035
-    broken = pythonAtLeast "3.12";
   };
 }