tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 pythonAtLeast,
7
8 # build-system
9 pybind11,
10 setuptools,
11 setuptools-scm,
12
13 # nativeBuildInputs
14 cmake,
15 ninja,
16
17 # dependencies
18 cloudpickle,
19 importlib-metadata,
20 numpy,
21 orjson,
22 packaging,
23 pyvers,
24 torch,
25
26 # tests
27 h5py,
28 pytestCheckHook,
29}:
30
31buildPythonPackage (finalAttrs: {
32 pname = "tensordict";
33 version = "0.11.0";
34 pyproject = true;
35
36 src = fetchFromGitHub {
37 owner = "pytorch";
38 repo = "tensordict";
39 tag = "v${finalAttrs.version}";
40 hash = "sha256-PUPDKv10Ks4B1kpgbRcnmfWFUkpFEdxMmTNztFVfdK4=";
41 };
42
43 postPatch = ''
44 substituteInPlace pyproject.toml \
45 --replace-fail "pybind11[global]" "pybind11"
46 '';
47
48 build-system = [
49 pybind11
50 setuptools
51 setuptools-scm
52 ];
53
54 nativeBuildInputs = [
55 cmake
56 ninja
57 ];
58 dontUseCmakeConfigure = true;
59
60 dependencies = [
61 cloudpickle
62 importlib-metadata
63 numpy
64 orjson
65 packaging
66 pyvers
67 torch
68 ];
69
70 pythonImportsCheck = [ "tensordict" ];
71
72 # We have to delete the source because otherwise it is used instead of the installed package.
73 preCheck = ''
74 rm -rf tensordict
75 '';
76
77 nativeCheckInputs = [
78 h5py
79 pytestCheckHook
80 ];
81
82 disabledTests = [
83 # FileNotFoundError: [Errno 2] No such file or directory: '/build/source/tensordict/tensorclass.pyi
84 "test_tensorclass_instance_methods"
85 "test_tensorclass_stub_methods"
86
87 # hangs forever on some CPUs
88 "test_map_iter_interrupt_early"
89 ]
90 ++ lib.optionals (pythonAtLeast "3.14") [
91 # AssertionError: assert 'a string!' == 'a metadata!'
92 "test_save_load_memmap"
93 ]
94 ++ lib.optionals stdenv.hostPlatform.isDarwin [
95 # Hangs due to the use of a pool
96 "test_chunksize_num_chunks"
97 "test_index_with_generator"
98 "test_map_exception"
99 "test_map"
100 "test_multiprocessing"
101 ];
102
103 disabledTestPaths = [
104 # torch._dynamo.exc.Unsupported: Graph break due to unsupported builtin None.ReferenceType.__new__.
105 "test/test_compile.py"
106 ]
107 ++ lib.optionals stdenv.hostPlatform.isDarwin [
108 # Hangs forever
109 "test/test_distributed.py"
110 # Hangs after testing due to pool usage
111 "test/test_h5.py"
112 "test/test_memmap.py"
113 ];
114
115 meta = {
116 description = "Pytorch dedicated tensor container";
117 changelog = "https://github.com/pytorch/tensordict/releases/tag/${finalAttrs.src.tag}";
118 homepage = "https://github.com/pytorch/tensordict";
119 license = lib.licenses.mit;
120 maintainers = with lib.maintainers; [ GaetanLepage ];
121 };
122})