Skip to content

[Distributed] Preserve tensorclass type across the wire - #1603

Merged
vmoens merged 3 commits into
gh/vmoens/81/basefrom
gh/vmoens/81/head
Feb 19, 2026
Merged

[Distributed] Preserve tensorclass type across the wire#1603
vmoens merged 3 commits into
gh/vmoens/81/basefrom
gh/vmoens/81/head

Conversation

@vmoens

@vmoens vmoens commented Feb 19, 2026

Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

Send the fully qualified tensorclass type string (e.g.
tensordict.testing.MyDistData) as metadata during init_remote and
broadcast. On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.

  • Add _tensorclass_type private param to init_remote / broadcast
    in TensorDictBase, injected into consolidation metadata.
  • Add _resolve_tensorclass_type helper to import a class from its
    dotted path.
  • Override init_remote, broadcast, from_remote_init on
    tensorclass to inject / recover the type automatically.
  • Create tensordict/testing.py with MyDistData so both sides of
    a multiprocessing test can import the same class.
  • Update all tensorclass distributed tests to assert type(result) is MyDistData.

Co-authored-by: Cursor cursoragent@cursor.com

[ghstack-poisoned]
vmoens added a commit that referenced this pull request Feb 19, 2026
Send the fully qualified tensorclass type string (e.g.
``tensordict.testing.MyDistData``) as metadata during init_remote and
broadcast.  On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.

- Add ``_tensorclass_type`` private param to ``init_remote`` / ``broadcast``
  in TensorDictBase, injected into consolidation metadata.
- Add ``_resolve_tensorclass_type`` helper to import a class from its
  dotted path.
- Override ``init_remote``, ``broadcast``, ``from_remote_init`` on
  tensorclass to inject / recover the type automatically.
- Create ``tensordict/testing.py`` with ``MyDistData`` so both sides of
  a multiprocessing test can import the same class.
- Update all tensorclass distributed tests to assert ``type(result) is
  MyDistData``.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: a7bf7c7
Pull-Request: #1603
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 19, 2026
@github-actions

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [Distributed].

Current title: [Distributed] Preserve tensorclass type across the wire

Supported Prefixes

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[BugFix] BugFix [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Perf [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] Setup [Setup] Update build configuration

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

@github-actions

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [Distributed].

Current title: [Distributed] Preserve tensorclass type across the wire

Supported Prefixes

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[BugFix] BugFix [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Perf [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] Setup [Setup] Update build configuration

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[ghstack-poisoned]
vmoens added a commit that referenced this pull request Feb 19, 2026
Send the fully qualified tensorclass type string (e.g.
``tensordict.testing.MyDistData``) as metadata during init_remote and
broadcast.  On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.

- Add ``_tensorclass_type`` private param to ``init_remote`` / ``broadcast``
  in TensorDictBase, injected into consolidation metadata.
- Add ``_resolve_tensorclass_type`` helper to import a class from its
  dotted path.
- Override ``init_remote``, ``broadcast``, ``from_remote_init`` on
  tensorclass to inject / recover the type automatically.
- Create ``tensordict/testing.py`` with ``MyDistData`` so both sides of
  a multiprocessing test can import the same class.
- Update all tensorclass distributed tests to assert ``type(result) is
  MyDistData``.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: 6f0af7f
Pull-Request: #1603
@github-actions

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [Distributed].

Current title: [Distributed] Preserve tensorclass type across the wire

Supported Prefixes

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[BugFix] BugFix [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Perf [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] Setup [Setup] Update build configuration

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

@github-actions

github-actions Bot commented Feb 19, 2026

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large&amp;#x26A0;\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 247. Improved: $\large\color{#35bf28}17$. Worsened: $\large\color{#d91a1a}9$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_plain_set_nested 38.5010μs 14.8395μs 67.3876 KOps/s 67.9850 KOps/s $\color{#d91a1a}-0.88\%$
test_plain_set_stack_nested 40.1310μs 15.1518μs 65.9987 KOps/s 66.7656 KOps/s $\color{#d91a1a}-1.15\%$
test_plain_set_nested_inplace 45.9310μs 16.3514μs 61.1570 KOps/s 61.3108 KOps/s $\color{#d91a1a}-0.25\%$
test_plain_set_stack_nested_inplace 49.7210μs 16.2885μs 61.3929 KOps/s 60.8887 KOps/s $\color{#35bf28}+0.83\%$
test_items 58.5610μs 6.0112μs 166.3559 KOps/s 174.5380 KOps/s $\color{#d91a1a}-4.69\%$
test_items_nested 0.5256ms 0.4579ms 2.1839 KOps/s 2.1768 KOps/s $\color{#35bf28}+0.32\%$
test_items_nested_locked 0.6330ms 0.4551ms 2.1975 KOps/s 2.1464 KOps/s $\color{#35bf28}+2.38\%$
test_items_nested_leaf 0.1370ms 97.0545μs 10.3035 KOps/s 10.2807 KOps/s $\color{#35bf28}+0.22\%$
test_items_stack_nested 0.5135ms 0.4526ms 2.2092 KOps/s 2.1923 KOps/s $\color{#35bf28}+0.77\%$
test_items_stack_nested_leaf 0.1311ms 95.8592μs 10.4320 KOps/s 10.2575 KOps/s $\color{#35bf28}+1.70\%$
test_items_stack_nested_locked 0.5177ms 0.4587ms 2.1799 KOps/s 2.1563 KOps/s $\color{#35bf28}+1.09\%$
test_keys 28.6710μs 4.1959μs 238.3253 KOps/s 237.3137 KOps/s $\color{#35bf28}+0.43\%$
test_keys_nested 0.1700ms 0.1207ms 8.2851 KOps/s 8.3127 KOps/s $\color{#d91a1a}-0.33\%$
test_keys_nested_locked 2.2105ms 0.1299ms 7.6986 KOps/s 7.0529 KOps/s $\textbf{\color{#35bf28}+9.15\%}$
test_keys_nested_leaf 0.1699ms 0.1109ms 9.0194 KOps/s 8.9526 KOps/s $\color{#35bf28}+0.75\%$
test_keys_stack_nested 0.1733ms 0.1210ms 8.2660 KOps/s 8.2949 KOps/s $\color{#d91a1a}-0.35\%$
test_keys_stack_nested_leaf 0.1471ms 0.1120ms 8.9319 KOps/s 9.0503 KOps/s $\color{#d91a1a}-1.31\%$
test_keys_stack_nested_locked 0.1646ms 0.1285ms 7.7806 KOps/s 7.7533 KOps/s $\color{#35bf28}+0.35\%$
test_values 6.1260μs 1.0330μs 968.0366 KOps/s 986.4791 KOps/s $\color{#d91a1a}-1.87\%$
test_values_nested 77.4320μs 47.5635μs 21.0245 KOps/s 21.0950 KOps/s $\color{#d91a1a}-0.33\%$
test_values_nested_locked 84.8310μs 50.9576μs 19.6241 KOps/s 19.6824 KOps/s $\color{#d91a1a}-0.30\%$
test_values_nested_leaf 86.7620μs 53.9908μs 18.5217 KOps/s 18.4441 KOps/s $\color{#35bf28}+0.42\%$
test_values_stack_nested 71.6810μs 47.5573μs 21.0273 KOps/s 20.9027 KOps/s $\color{#35bf28}+0.60\%$
test_values_stack_nested_leaf 93.7020μs 54.6911μs 18.2845 KOps/s 18.4162 KOps/s $\color{#d91a1a}-0.72\%$
test_values_stack_nested_locked 1.7498ms 50.9647μs 19.6214 KOps/s 19.5958 KOps/s $\color{#35bf28}+0.13\%$
test_membership 5.0250μs 0.8481μs 1.1791 MOps/s 1.1889 MOps/s $\color{#d91a1a}-0.82\%$
test_membership_nested 0.6503ms 2.8948μs 345.4456 KOps/s 342.9419 KOps/s $\color{#35bf28}+0.73\%$
test_membership_nested_leaf 27.9000μs 2.9298μs 341.3186 KOps/s 354.8948 KOps/s $\color{#d91a1a}-3.83\%$
test_membership_stacked_nested 26.2610μs 2.9147μs 343.0884 KOps/s 343.6985 KOps/s $\color{#d91a1a}-0.18\%$
test_membership_stacked_nested_leaf 26.1110μs 2.9528μs 338.6633 KOps/s 343.8696 KOps/s $\color{#d91a1a}-1.51\%$
test_membership_nested_last 40.0400μs 4.3488μs 229.9477 KOps/s 228.6212 KOps/s $\color{#35bf28}+0.58\%$
test_membership_nested_leaf_last 38.4010μs 4.3206μs 231.4512 KOps/s 229.1702 KOps/s $\color{#35bf28}+1.00\%$
test_membership_stacked_nested_last 38.1400μs 4.3789μs 228.3695 KOps/s 232.2286 KOps/s $\color{#d91a1a}-1.66\%$
test_membership_stacked_nested_leaf_last 61.2110μs 4.3363μs 230.6127 KOps/s 232.9536 KOps/s $\color{#d91a1a}-1.00\%$
test_nested_getleaf 50.0510μs 21.7073μs 46.0675 KOps/s 47.9274 KOps/s $\color{#d91a1a}-3.88\%$
test_nested_get 73.5210μs 20.5616μs 48.6344 KOps/s 50.5474 KOps/s $\color{#d91a1a}-3.78\%$
test_stacked_getleaf 55.3510μs 21.6017μs 46.2927 KOps/s 47.8468 KOps/s $\color{#d91a1a}-3.25\%$
test_stacked_get 46.7210μs 20.4593μs 48.8775 KOps/s 50.3856 KOps/s $\color{#d91a1a}-2.99\%$
test_nested_getitemleaf 53.6810μs 22.3751μs 44.6925 KOps/s 47.3050 KOps/s $\textbf{\color{#d91a1a}-5.52\%}$
test_nested_getitem 48.7200μs 21.1130μs 47.3641 KOps/s 49.6239 KOps/s $\color{#d91a1a}-4.55\%$
test_stacked_getitemleaf 0.1054ms 21.9232μs 45.6137 KOps/s 47.0451 KOps/s $\color{#d91a1a}-3.04\%$
test_stacked_getitem 47.1710μs 21.0956μs 47.4033 KOps/s 49.7503 KOps/s $\color{#d91a1a}-4.72\%$
test_lock_nested 0.7240ms 0.4716ms 2.1203 KOps/s 2.0805 KOps/s $\color{#35bf28}+1.91\%$
test_lock_stack_nested 0.5342ms 0.4780ms 2.0918 KOps/s 2.0791 KOps/s $\color{#35bf28}+0.61\%$
test_unlock_nested 0.6034ms 0.3805ms 2.6279 KOps/s 2.6185 KOps/s $\color{#35bf28}+0.36\%$
test_unlock_stack_nested 0.4358ms 0.3841ms 2.6032 KOps/s 2.5808 KOps/s $\color{#35bf28}+0.87\%$
test_flatten_speed 0.1884ms 0.1238ms 8.0781 KOps/s 8.1550 KOps/s $\color{#d91a1a}-0.94\%$
test_unflatten_speed 0.6158ms 0.5687ms 1.7585 KOps/s 1.7436 KOps/s $\color{#35bf28}+0.86\%$
test_common_ops 0.8269ms 0.6910ms 1.4472 KOps/s 1.4474 KOps/s $\color{#d91a1a}-0.01\%$
test_creation 67.9110μs 2.9366μs 340.5277 KOps/s 340.5763 KOps/s $\color{#d91a1a}-0.01\%$
test_creation_empty 37.3600μs 6.1475μs 162.6690 KOps/s 159.3918 KOps/s $\color{#35bf28}+2.06\%$
test_creation_nested_1 39.6210μs 10.4552μs 95.6458 KOps/s 95.2344 KOps/s $\color{#35bf28}+0.43\%$
test_creation_nested_2 43.5910μs 11.8577μs 84.3332 KOps/s 84.0492 KOps/s $\color{#35bf28}+0.34\%$
test_creation_many_keys[10] 46.6110μs 18.2447μs 54.8105 KOps/s 54.0859 KOps/s $\color{#35bf28}+1.34\%$
test_creation_many_keys[50] 0.1081ms 77.9439μs 12.8297 KOps/s 12.7910 KOps/s $\color{#35bf28}+0.30\%$
test_creation_many_keys[100] 0.1963ms 0.1525ms 6.5570 KOps/s 6.5386 KOps/s $\color{#35bf28}+0.28\%$
test_creation_nested_many_keys[10] 71.9810μs 39.1418μs 25.5482 KOps/s 25.5653 KOps/s $\color{#d91a1a}-0.07\%$
test_creation_nested_many_keys[50] 0.2011ms 0.1609ms 6.2154 KOps/s 6.2276 KOps/s $\color{#d91a1a}-0.20\%$
test_clone 45.0510μs 13.5765μs 73.6565 KOps/s 73.6061 KOps/s $\color{#35bf28}+0.07\%$
test_getitem[int] 1.7113ms 14.6250μs 68.3759 KOps/s 57.1400 KOps/s $\textbf{\color{#35bf28}+19.66\%}$
test_getitem[slice_int] 0.1444ms 26.6540μs 37.5178 KOps/s 40.6287 KOps/s $\textbf{\color{#d91a1a}-7.66\%}$
test_getitem[range] 0.1748ms 62.8313μs 15.9156 KOps/s 15.3636 KOps/s $\color{#35bf28}+3.59\%$
test_getitem[tuple] 0.1488ms 23.9813μs 41.6991 KOps/s 42.2646 KOps/s $\color{#d91a1a}-1.34\%$
test_getitem[list] 0.1890ms 61.0387μs 16.3831 KOps/s 16.5656 KOps/s $\color{#d91a1a}-1.10\%$
test_setitem_dim[int] 48.4210μs 27.8113μs 35.9567 KOps/s 36.2696 KOps/s $\color{#d91a1a}-0.86\%$
test_setitem_dim[slice_int] 69.6610μs 44.9166μs 22.2635 KOps/s 22.5365 KOps/s $\color{#d91a1a}-1.21\%$
test_setitem_dim[range] 0.1276ms 98.3358μs 10.1692 KOps/s 10.0726 KOps/s $\color{#35bf28}+0.96\%$
test_setitem_dim[tuple] 65.1310μs 42.7804μs 23.3752 KOps/s 23.3336 KOps/s $\color{#35bf28}+0.18\%$
test_setitem 71.6220μs 17.8289μs 56.0887 KOps/s 55.0444 KOps/s $\color{#35bf28}+1.90\%$
test_set 66.9520μs 17.1116μs 58.4399 KOps/s 57.9852 KOps/s $\color{#35bf28}+0.78\%$
test_set_shared 0.4865ms 0.2045ms 4.8903 KOps/s 4.7073 KOps/s $\color{#35bf28}+3.89\%$
test_update 0.3234ms 21.5215μs 46.4651 KOps/s 43.3864 KOps/s $\textbf{\color{#35bf28}+7.10\%}$
test_update_nested 64.8610μs 33.4952μs 29.8550 KOps/s 28.5648 KOps/s $\color{#35bf28}+4.52\%$
test_update__nested 0.4641ms 34.4249μs 29.0487 KOps/s 28.4719 KOps/s $\color{#35bf28}+2.03\%$
test_set_nested 60.4620μs 19.2294μs 52.0038 KOps/s 51.2241 KOps/s $\color{#35bf28}+1.52\%$
test_set_nested_new 63.4920μs 23.7763μs 42.0587 KOps/s 38.6229 KOps/s $\textbf{\color{#35bf28}+8.90\%}$
test_select 85.0820μs 41.5522μs 24.0661 KOps/s 23.4400 KOps/s $\color{#35bf28}+2.67\%$
test_select_nested 0.1087ms 74.8673μs 13.3570 KOps/s 13.5005 KOps/s $\color{#d91a1a}-1.06\%$
test_exclude_nested 0.1274ms 91.9732μs 10.8727 KOps/s 10.9001 KOps/s $\color{#d91a1a}-0.25\%$
test_empty[True] 0.4471ms 0.3860ms 2.5904 KOps/s 2.5448 KOps/s $\color{#35bf28}+1.79\%$
test_empty[False] 9.9277μs 1.3281μs 752.9502 KOps/s 753.1475 KOps/s $\color{#d91a1a}-0.03\%$
test_to 0.1031ms 71.0315μs 14.0783 KOps/s 12.6412 KOps/s $\textbf{\color{#35bf28}+11.37\%}$
test_to_nonblocking 0.1075ms 66.7184μs 14.9884 KOps/s 14.6485 KOps/s $\color{#35bf28}+2.32\%$
test_unbind_speed 0.3718ms 0.3274ms 3.0546 KOps/s 3.0669 KOps/s $\color{#d91a1a}-0.40\%$
test_unbind_speed_stack0 0.3877ms 0.3303ms 3.0274 KOps/s 3.0798 KOps/s $\color{#d91a1a}-1.70\%$
test_unbind_speed_stack1 0.1042s 0.8337ms 1.1995 KOps/s 1.1888 KOps/s $\color{#35bf28}+0.90\%$
test_split 0.1037s 1.2702ms 787.2792 Ops/s 786.1097 Ops/s $\color{#35bf28}+0.15\%$
test_chunk 0.1041s 1.2113ms 825.5680 Ops/s 924.9371 Ops/s $\textbf{\color{#d91a1a}-10.74\%}$
test_to_cpu_blocking 29.1362ms 28.6235ms 34.9363 Ops/s 40.0692 Ops/s $\textbf{\color{#d91a1a}-12.81\%}$
test_to_cpu_global_sync 11.6288ms 11.2858ms 88.6070 Ops/s 89.4632 Ops/s $\color{#d91a1a}-0.96\%$
test_to_cpu_event_sync 12.5525ms 12.2533ms 81.6107 Ops/s 82.3870 Ops/s $\color{#d91a1a}-0.94\%$
test_to_cpu_default 0.1163s 13.5734ms 73.6734 Ops/s 82.2807 Ops/s $\textbf{\color{#d91a1a}-10.46\%}$
test_consolidate[False-None] 4.2453ms 4.1153ms 242.9936 Ops/s 219.6502 Ops/s $\textbf{\color{#35bf28}+10.63\%}$
test_consolidate[default-None] 2.1139ms 2.0284ms 493.0039 Ops/s 483.6884 Ops/s $\color{#35bf28}+1.93\%$
test_consolidate[reduce-overhead-None] 2.0521ms 1.9484ms 513.2524 Ops/s 511.2236 Ops/s $\color{#35bf28}+0.40\%$
test_consolidate_njt[False-None] 8.6808ms 8.4068ms 118.9509 Ops/s 118.6549 Ops/s $\color{#35bf28}+0.25\%$
test_to[False-False-None] 2.2347ms 2.0668ms 483.8292 Ops/s 476.1794 Ops/s $\color{#35bf28}+1.61\%$
test_to[True-False-None] 2.1857ms 1.9143ms 522.3787 Ops/s 513.1205 Ops/s $\color{#35bf28}+1.80\%$
test_to[within-False-None] 6.2058ms 6.1043ms 163.8196 Ops/s 163.1102 Ops/s $\color{#35bf28}+0.43\%$
test_to[True-default-None] 7.8129ms 7.5171ms 133.0295 Ops/s 129.6589 Ops/s $\color{#35bf28}+2.60\%$
test_to_njt[False-False-None] 8.6629ms 8.5187ms 117.3883 Ops/s 115.9740 Ops/s $\color{#35bf28}+1.22\%$
test_to_njt[True-False-None] 7.1268ms 6.9084ms 144.7511 Ops/s 140.0232 Ops/s $\color{#35bf28}+3.38\%$
test_to_njt[within-False-None] 15.6610ms 15.5015ms 64.5098 Ops/s 62.9410 Ops/s $\color{#35bf28}+2.49\%$
test_creation[device0] 0.4028ms 0.1167ms 8.5662 KOps/s 8.3004 KOps/s $\color{#35bf28}+3.20\%$
test_creation_from_tensor 0.4546ms 0.1141ms 8.7660 KOps/s 8.5158 KOps/s $\color{#35bf28}+2.94\%$
test_add_one[memmap_tensor0] 0.3724ms 6.5465μs 152.7544 KOps/s 147.2731 KOps/s $\color{#35bf28}+3.72\%$
test_contiguous[memmap_tensor0] 26.8100μs 0.6658μs 1.5020 MOps/s 2.1737 MOps/s $\textbf{\color{#d91a1a}-30.90\%}$
test_stack[memmap_tensor0] 20.4000μs 4.5916μs 217.7896 KOps/s 219.3409 KOps/s $\color{#d91a1a}-0.71\%$
test_memmaptd_index 1.0186ms 0.2662ms 3.7572 KOps/s 3.8583 KOps/s $\color{#d91a1a}-2.62\%$
test_memmaptd_index_astensor 0.5458ms 0.3592ms 2.7842 KOps/s 2.8598 KOps/s $\color{#d91a1a}-2.64\%$
test_memmaptd_index_op 0.8432ms 0.6004ms 1.6656 KOps/s 1.6600 KOps/s $\color{#35bf28}+0.34\%$
test_serialize_model 0.3185s 0.1630s 6.1361 Ops/s 7.3760 Ops/s $\textbf{\color{#d91a1a}-16.81\%}$
test_serialize_model_pickle 1.3498s 1.1920s 0.8389 Ops/s 0.8266 Ops/s $\color{#35bf28}+1.48\%$
test_serialize_weights 0.1397s 0.1363s 7.3389 Ops/s 7.3675 Ops/s $\color{#d91a1a}-0.39\%$
test_serialize_weights_returnearly 0.4557s 94.2220ms 10.6132 Ops/s 5.9223 Ops/s $\textbf{\color{#35bf28}+79.21\%}$
test_serialize_weights_pickle 1.3647s 1.2132s 0.8242 Ops/s 0.8238 Ops/s $\color{#35bf28}+0.05\%$
test_reshape_pytree 0.2153ms 33.6137μs 29.7498 KOps/s 29.9105 KOps/s $\color{#d91a1a}-0.54\%$
test_reshape_td 76.8920μs 44.8032μs 22.3198 KOps/s 22.7702 KOps/s $\color{#d91a1a}-1.98\%$
test_view_pytree 0.2336ms 33.3024μs 30.0279 KOps/s 30.9930 KOps/s $\color{#d91a1a}-3.11\%$
test_view_td 96.0510μs 51.9268μs 19.2579 KOps/s 18.7643 KOps/s $\color{#35bf28}+2.63\%$
test_unbind_pytree 0.2409ms 36.7415μs 27.2172 KOps/s 27.0357 KOps/s $\color{#35bf28}+0.67\%$
test_unbind_td 0.1379ms 48.6993μs 20.5342 KOps/s 20.3478 KOps/s $\color{#35bf28}+0.92\%$
test_split_pytree 0.2489ms 42.9948μs 23.2586 KOps/s 23.7785 KOps/s $\color{#d91a1a}-2.19\%$
test_split_td 0.2076ms 65.5165μs 15.2633 KOps/s 15.3953 KOps/s $\color{#d91a1a}-0.86\%$
test_add_pytree 0.1960ms 42.3189μs 23.6301 KOps/s 23.7199 KOps/s $\color{#d91a1a}-0.38\%$
test_add_td 93.8920μs 53.1536μs 18.8134 KOps/s 18.2957 KOps/s $\color{#35bf28}+2.83\%$
test_compile_add_one_nested[tensordict-compile] 0.1940ms 0.1416ms 7.0626 KOps/s 6.4524 KOps/s $\textbf{\color{#35bf28}+9.46\%}$
test_compile_add_one_nested[tensordict-eager] 0.3121ms 0.1951ms 5.1269 KOps/s 5.1849 KOps/s $\color{#d91a1a}-1.12\%$
test_compile_add_one_nested[pytree-compile] 0.1584ms 0.1095ms 9.1306 KOps/s 8.9889 KOps/s $\color{#35bf28}+1.58\%$
test_compile_add_one_nested[pytree-eager] 0.4469ms 0.1821ms 5.4912 KOps/s 5.3821 KOps/s $\color{#35bf28}+2.03\%$
test_compile_copy_nested[tensordict-compile] 76.2410μs 29.6336μs 33.7454 KOps/s 30.4929 KOps/s $\textbf{\color{#35bf28}+10.67\%}$
test_compile_copy_nested[tensordict-eager] 84.9820μs 52.9126μs 18.8991 KOps/s 18.9621 KOps/s $\color{#d91a1a}-0.33\%$
test_compile_copy_nested[pytree-compile] 39.8410μs 9.9477μs 100.5257 KOps/s 101.8539 KOps/s $\color{#d91a1a}-1.30\%$
test_compile_copy_nested[pytree-eager] 0.4615ms 70.4838μs 14.1877 KOps/s 14.6012 KOps/s $\color{#d91a1a}-2.83\%$
test_compile_add_one_flat[tensordict-compile] 0.2301ms 0.1786ms 5.5994 KOps/s 5.3595 KOps/s $\color{#35bf28}+4.48\%$
test_compile_add_one_flat[tensordict-eager] 6.2944ms 0.2579ms 3.8767 KOps/s 3.8938 KOps/s $\color{#d91a1a}-0.44\%$
test_compile_add_one_flat[tensorclass-compile] 0.1557ms 0.1179ms 8.4786 KOps/s 8.3221 KOps/s $\color{#35bf28}+1.88\%$
test_compile_add_one_flat[tensorclass-eager] 0.1055ms 68.9842μs 14.4961 KOps/s 14.3627 KOps/s $\color{#35bf28}+0.93\%$
test_compile_add_one_flat[pytree-compile] 0.1966ms 0.1595ms 6.2701 KOps/s 6.1121 KOps/s $\color{#35bf28}+2.59\%$
test_compile_add_one_flat[pytree-eager] 0.8337ms 0.5342ms 1.8720 KOps/s 1.7855 KOps/s $\color{#35bf28}+4.84\%$
test_compile_add_self_flat[tensordict-eager] 0.3638ms 0.3104ms 3.2218 KOps/s 3.2077 KOps/s $\color{#35bf28}+0.44\%$
test_compile_add_self_flat[tensordict-compile] 0.2342ms 0.1804ms 5.5441 KOps/s 4.9742 KOps/s $\textbf{\color{#35bf28}+11.46\%}$
test_compile_add_self_flat[tensorclass-eager] 0.1250ms 84.7258μs 11.8028 KOps/s 11.7279 KOps/s $\color{#35bf28}+0.64\%$
test_compile_add_self_flat[tensorclass-compile] 0.1590ms 0.1211ms 8.2591 KOps/s 7.9143 KOps/s $\color{#35bf28}+4.36\%$
test_compile_add_self_flat[pytree-eager] 0.6800ms 0.4417ms 2.2642 KOps/s 2.1575 KOps/s $\color{#35bf28}+4.94\%$
test_compile_add_self_flat[pytree-compile] 0.2026ms 0.1594ms 6.2729 KOps/s 5.7894 KOps/s $\textbf{\color{#35bf28}+8.35\%}$
test_compile_copy_flat[tensordict-compile] 60.1010μs 24.2722μs 41.1993 KOps/s 38.3630 KOps/s $\textbf{\color{#35bf28}+7.39\%}$
test_compile_copy_flat[tensordict-eager] 69.1010μs 41.9963μs 23.8116 KOps/s 24.0958 KOps/s $\color{#d91a1a}-1.18\%$
test_compile_copy_flat[pytree-compile] 46.7610μs 10.8535μs 92.1362 KOps/s 92.4748 KOps/s $\color{#d91a1a}-0.37\%$
test_compile_copy_flat[pytree-eager] 0.1807s 62.0356μs 16.1198 KOps/s 19.2389 KOps/s $\textbf{\color{#d91a1a}-16.21\%}$
test_compile_assign_and_add[tensordict-compile] 2.0426ms 0.1798ms 5.5621 KOps/s 5.2644 KOps/s $\textbf{\color{#35bf28}+5.65\%}$
test_compile_assign_and_add[tensordict-eager] 3.4550ms 3.2978ms 303.2331 Ops/s 301.3114 Ops/s $\color{#35bf28}+0.64\%$
test_compile_assign_and_add[pytree-compile] 1.9902ms 0.1625ms 6.1531 KOps/s 6.1412 KOps/s $\color{#35bf28}+0.19\%$
test_compile_assign_and_add[pytree-eager] 2.9409ms 2.8168ms 355.0068 Ops/s 350.6682 Ops/s $\color{#35bf28}+1.24\%$
test_compile_indexing[tensor-tensordict-compile] 0.1712ms 0.1096ms 9.1221 KOps/s 8.9213 KOps/s $\color{#35bf28}+2.25\%$
test_compile_indexing[tensor-tensordict-eager] 0.3651ms 72.5013μs 13.7929 KOps/s 13.8126 KOps/s $\color{#d91a1a}-0.14\%$
test_compile_indexing[tensor-tensorclass-compile] 0.1368ms 97.4024μs 10.2667 KOps/s 10.2187 KOps/s $\color{#35bf28}+0.47\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2708ms 44.4807μs 22.4817 KOps/s 22.2434 KOps/s $\color{#35bf28}+1.07\%$
test_compile_indexing[tensor-pytree-compile] 0.1452ms 98.0281μs 10.2012 KOps/s 10.1776 KOps/s $\color{#35bf28}+0.23\%$
test_compile_indexing[tensor-pytree-eager] 0.2597ms 44.6191μs 22.4119 KOps/s 22.2359 KOps/s $\color{#35bf28}+0.79\%$
test_compile_indexing[slice-tensordict-compile] 0.1279ms 56.0741μs 17.8336 KOps/s 16.9567 KOps/s $\textbf{\color{#35bf28}+5.17\%}$
test_compile_indexing[slice-tensordict-eager] 0.2254ms 27.6765μs 36.1317 KOps/s 36.3065 KOps/s $\color{#d91a1a}-0.48\%$
test_compile_indexing[slice-tensorclass-compile] 82.4520μs 45.9461μs 21.7646 KOps/s 22.1038 KOps/s $\color{#d91a1a}-1.53\%$
test_compile_indexing[slice-tensorclass-eager] 0.2601ms 22.8239μs 43.8137 KOps/s 43.9644 KOps/s $\color{#d91a1a}-0.34\%$
test_compile_indexing[slice-pytree-compile] 86.9220μs 46.0157μs 21.7317 KOps/s 21.8590 KOps/s $\color{#d91a1a}-0.58\%$
test_compile_indexing[slice-pytree-eager] 0.2696ms 22.6049μs 44.2381 KOps/s 44.2704 KOps/s $\color{#d91a1a}-0.07\%$
test_compile_indexing[int-tensordict-compile] 0.1073ms 56.9620μs 17.5556 KOps/s 16.8742 KOps/s $\color{#35bf28}+4.04\%$
test_compile_indexing[int-tensordict-eager] 0.2304ms 27.5857μs 36.2507 KOps/s 36.1035 KOps/s $\color{#35bf28}+0.41\%$
test_compile_indexing[int-tensorclass-compile] 93.0420μs 45.4771μs 21.9891 KOps/s 21.7999 KOps/s $\color{#35bf28}+0.87\%$
test_compile_indexing[int-tensorclass-eager] 0.2617ms 22.7833μs 43.8918 KOps/s 44.1711 KOps/s $\color{#d91a1a}-0.63\%$
test_compile_indexing[int-pytree-compile] 92.7320μs 45.6473μs 21.9071 KOps/s 21.7277 KOps/s $\color{#35bf28}+0.83\%$
test_compile_indexing[int-pytree-eager] 0.2611ms 22.1474μs 45.1520 KOps/s 44.1670 KOps/s $\color{#35bf28}+2.23\%$
test_mod_add[eager] 91.3210μs 48.8463μs 20.4724 KOps/s 19.9961 KOps/s $\color{#35bf28}+2.38\%$
test_mod_add[compile] 0.2098ms 0.1055ms 9.4815 KOps/s 9.4581 KOps/s $\color{#35bf28}+0.25\%$
test_mod_add[compile-overhead] 0.2355ms 0.1499ms 6.6695 KOps/s 6.5438 KOps/s $\color{#35bf28}+1.92\%$
test_mod_wrap[eager] 0.3616ms 0.2901ms 3.4475 KOps/s 3.3618 KOps/s $\color{#35bf28}+2.55\%$
test_mod_wrap[compile] 0.4528ms 0.3612ms 2.7689 KOps/s 2.8293 KOps/s $\color{#d91a1a}-2.14\%$
test_mod_wrap[compile-overhead] 7.1843ms 3.9842ms 250.9898 Ops/s 247.7151 Ops/s $\color{#35bf28}+1.32\%$
test_mod_wrap_and_backward[eager] 1.6524ms 1.5273ms 654.7638 Ops/s 646.8489 Ops/s $\color{#35bf28}+1.22\%$
test_mod_wrap_and_backward[compile] 1.6436ms 1.4504ms 689.4617 Ops/s 686.4058 Ops/s $\color{#35bf28}+0.45\%$
test_mod_wrap_and_backward[compile-overhead] 1.2579ms 0.9570ms 1.0449 KOps/s 1.1021 KOps/s $\textbf{\color{#d91a1a}-5.19\%}$
test_seq_add[eager] 0.2586ms 0.1539ms 6.4984 KOps/s 6.3409 KOps/s $\color{#35bf28}+2.48\%$
test_seq_add[compile] 0.4629ms 0.1179ms 8.4791 KOps/s 8.3965 KOps/s $\color{#35bf28}+0.98\%$
test_seq_add[compile-overhead] 0.1987ms 0.1564ms 6.3945 KOps/s 6.2763 KOps/s $\color{#35bf28}+1.88\%$
test_seq_wrap[eager] 0.7947ms 0.5475ms 1.8265 KOps/s 1.8184 KOps/s $\color{#35bf28}+0.45\%$
test_seq_wrap[compile] 0.4159ms 0.3664ms 2.7295 KOps/s 2.6170 KOps/s $\color{#35bf28}+4.30\%$
test_seq_wrap[compile-overhead] 0.3425ms 0.2666ms 3.7511 KOps/s 3.7069 KOps/s $\color{#35bf28}+1.19\%$
test_func_call_runtime[False-eager] 0.9764ms 0.8846ms 1.1305 KOps/s 1.1597 KOps/s $\color{#d91a1a}-2.52\%$
test_func_call_runtime[False-compile] 1.0121ms 0.9092ms 1.0998 KOps/s 1.0833 KOps/s $\color{#35bf28}+1.52\%$
test_func_call_runtime[False-compile-overhead] 0.5319ms 0.4602ms 2.1728 KOps/s 2.1262 KOps/s $\color{#35bf28}+2.19\%$
test_func_call_runtime[True-eager] 1.1925ms 1.0922ms 915.6017 Ops/s 908.6231 Ops/s $\color{#35bf28}+0.77\%$
test_func_call_runtime[True-compile] 0.9852ms 0.9197ms 1.0873 KOps/s 1.0646 KOps/s $\color{#35bf28}+2.13\%$
test_func_call_runtime[True-compile-overhead] 0.5237ms 0.4737ms 2.1112 KOps/s 2.0567 KOps/s $\color{#35bf28}+2.65\%$
test_func_call_cm_runtime[False-eager] 0.9136ms 0.8388ms 1.1922 KOps/s 1.1121 KOps/s $\textbf{\color{#35bf28}+7.20\%}$
test_func_call_cm_runtime[False-compile] 1.1313ms 0.9136ms 1.0945 KOps/s 1.0707 KOps/s $\color{#35bf28}+2.22\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5232ms 0.4635ms 2.1575 KOps/s 2.1150 KOps/s $\color{#35bf28}+2.01\%$
test_func_call_cm_runtime[True-eager] 1.4306ms 1.2332ms 810.9048 Ops/s 803.0255 Ops/s $\color{#35bf28}+0.98\%$
test_func_call_cm_runtime[True-compile] 1.2216ms 0.9662ms 1.0350 KOps/s 1.0307 KOps/s $\color{#35bf28}+0.42\%$
test_func_call_cm_runtime[True-compile-overhead] 0.5855ms 0.5088ms 1.9655 KOps/s 1.9280 KOps/s $\color{#35bf28}+1.95\%$
test_vmap_func_call_cm_runtime[eager] 2.8750ms 2.3656ms 422.7308 Ops/s 413.9631 Ops/s $\color{#35bf28}+2.12\%$
test_vmap_func_call_cm_runtime[compile] 1.1242ms 0.9833ms 1.0170 KOps/s 1.0007 KOps/s $\color{#35bf28}+1.62\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.6301ms 0.5132ms 1.9486 KOps/s 1.9061 KOps/s $\color{#35bf28}+2.23\%$
test_distributed 3.0330ms 0.1642ms 6.0890 KOps/s 6.0613 KOps/s $\color{#35bf28}+0.46\%$
test_tdmodule 68.8010μs 27.9745μs 35.7468 KOps/s 34.2219 KOps/s $\color{#35bf28}+4.46\%$
test_tdmodule_dispatch 76.4320μs 45.1153μs 22.1654 KOps/s 21.9300 KOps/s $\color{#35bf28}+1.07\%$
test_tdseq 47.7010μs 27.6821μs 36.1245 KOps/s 35.8260 KOps/s $\color{#35bf28}+0.83\%$
test_tdseq_dispatch 65.7120μs 47.6163μs 21.0012 KOps/s 20.8077 KOps/s $\color{#35bf28}+0.93\%$
test_instantiation_functorch 2.1974ms 2.0802ms 480.7224 Ops/s 480.0977 Ops/s $\color{#35bf28}+0.13\%$
test_exec_functorch 0.2506ms 0.1816ms 5.5078 KOps/s 5.4880 KOps/s $\color{#35bf28}+0.36\%$
test_exec_functional_call 0.2909ms 0.1682ms 5.9445 KOps/s 6.1415 KOps/s $\color{#d91a1a}-3.21\%$
test_exec_td_decorator 0.4441ms 0.2400ms 4.1667 KOps/s 4.2294 KOps/s $\color{#d91a1a}-1.48\%$
test_vmap_mlp_speed_decorator[True-True] 0.9957ms 0.8245ms 1.2129 KOps/s 1.2015 KOps/s $\color{#35bf28}+0.95\%$
test_vmap_mlp_speed_decorator[True-False] 1.0538ms 0.8235ms 1.2144 KOps/s 1.1984 KOps/s $\color{#35bf28}+1.34\%$
test_vmap_mlp_speed_decorator[False-True] 0.9233ms 0.7172ms 1.3943 KOps/s 1.3696 KOps/s $\color{#35bf28}+1.81\%$
test_vmap_mlp_speed_decorator[False-False] 0.9032ms 0.7131ms 1.4024 KOps/s 1.3705 KOps/s $\color{#35bf28}+2.33\%$
test_vmap_transformer_speed_decorator[True-True] 20.8543ms 20.6672ms 48.3858 Ops/s 47.8091 Ops/s $\color{#35bf28}+1.21\%$
test_vmap_transformer_speed_decorator[True-False] 20.9638ms 20.7190ms 48.2649 Ops/s 47.6908 Ops/s $\color{#35bf28}+1.20\%$
test_vmap_transformer_speed_decorator[False-True] 21.1460ms 20.4988ms 48.7833 Ops/s 48.3007 Ops/s $\color{#35bf28}+1.00\%$
test_vmap_transformer_speed_decorator[False-False] 21.3641ms 20.6266ms 48.4811 Ops/s 48.4059 Ops/s $\color{#35bf28}+0.16\%$
test_to_module_speed[True] 2.0357ms 1.4709ms 679.8604 Ops/s 676.8233 Ops/s $\color{#35bf28}+0.45\%$
test_to_module_speed[False] 2.0130ms 1.4513ms 689.0379 Ops/s 685.2291 Ops/s $\color{#35bf28}+0.56\%$
test_tc_init 68.3910μs 45.5391μs 21.9592 KOps/s 21.5674 KOps/s $\color{#35bf28}+1.82\%$
test_tc_init_tensor_only 35.6210μs 9.7951μs 102.0919 KOps/s 99.8962 KOps/s $\color{#35bf28}+2.20\%$
test_tc_init_nested 0.1407ms 90.0015μs 11.1109 KOps/s 11.0029 KOps/s $\color{#35bf28}+0.98\%$
test_tc_init_many_fields 84.5320μs 16.4160μs 60.9163 KOps/s 58.5543 KOps/s $\color{#35bf28}+4.03\%$
test_tc_first_layer_tensor 18.3700μs 1.8571μs 538.4666 KOps/s 532.0797 KOps/s $\color{#35bf28}+1.20\%$
test_tc_first_layer_tensor_only 5.4701μs 0.7447μs 1.3427 MOps/s 1.3328 MOps/s $\color{#35bf28}+0.74\%$
test_tc_first_layer_tensor_set 37.9010μs 3.9787μs 251.3369 KOps/s 252.6251 KOps/s $\color{#d91a1a}-0.51\%$
test_tc_first_layer_tensor_only_set 24.3710μs 3.2087μs 311.6560 KOps/s 309.7789 KOps/s $\color{#35bf28}+0.61\%$
test_tc_first_layer_nontensor 37.4200μs 6.2621μs 159.6907 KOps/s 157.8029 KOps/s $\color{#35bf28}+1.20\%$
test_tc_second_layer_tensor 34.6600μs 4.4392μs 225.2683 KOps/s 226.3248 KOps/s $\color{#d91a1a}-0.47\%$
test_tc_second_layer_nontensor 41.7610μs 8.8626μs 112.8342 KOps/s 110.9538 KOps/s $\color{#35bf28}+1.69\%$
test_unbind 0.2602s 16.3130ms 61.3008 Ops/s 63.2965 Ops/s $\color{#d91a1a}-3.15\%$
test_full_like 5.5392ms 4.3759ms 228.5231 Ops/s 228.7359 Ops/s $\color{#d91a1a}-0.09\%$
test_zeros_like 4.4978ms 4.3630ms 229.1978 Ops/s 228.4111 Ops/s $\color{#35bf28}+0.34\%$
test_ones_like 4.5006ms 4.3699ms 228.8381 Ops/s 228.8494 Ops/s $-0.00\%$
test_clone 6.5104ms 6.4108ms 155.9859 Ops/s 154.7485 Ops/s $\color{#35bf28}+0.80\%$
test_squeeze 0.2002ms 14.5616μs 68.6737 KOps/s 68.6058 KOps/s $\color{#35bf28}+0.10\%$
test_unsqueeze 0.2688ms 0.1101ms 9.0827 KOps/s 8.9751 KOps/s $\color{#35bf28}+1.20\%$
test_split 0.2373ms 0.1858ms 5.3819 KOps/s 5.3879 KOps/s $\color{#d91a1a}-0.11\%$
test_permute 0.3183ms 0.2111ms 4.7380 KOps/s 4.8883 KOps/s $\color{#d91a1a}-3.08\%$
test_stack 51.6150ms 50.8354ms 19.6713 Ops/s 19.5856 Ops/s $\color{#35bf28}+0.44\%$
test_cat 51.5173ms 51.2171ms 19.5247 Ops/s 19.5567 Ops/s $\color{#d91a1a}-0.16\%$
test_sequential_tensordict 0.3314ms 0.2206ms 4.5334 KOps/s 4.4131 KOps/s $\color{#35bf28}+2.73\%$
test_sequential_graph_module 0.5100ms 0.1170ms 8.5478 KOps/s 7.9694 KOps/s $\textbf{\color{#35bf28}+7.26\%}$
test_nested_tensordict 0.4158ms 0.2796ms 3.5761 KOps/s 3.4851 KOps/s $\color{#35bf28}+2.61\%$
test_nested_graph_module 0.5400ms 0.1264ms 7.9135 KOps/s 7.3327 KOps/s $\textbf{\color{#35bf28}+7.92\%}$

@github-actions

github-actions Bot commented Feb 19, 2026

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large&amp;#x26A0;\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 247. Improved: $\large\color{#35bf28}26$. Worsened: $\large\color{#d91a1a}5$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_plain_set_nested 34.1210μs 14.7592μs 67.7545 KOps/s 67.8203 KOps/s $\color{#d91a1a}-0.10\%$
test_plain_set_stack_nested 45.5910μs 14.9035μs 67.0981 KOps/s 67.2045 KOps/s $\color{#d91a1a}-0.16\%$
test_plain_set_nested_inplace 36.6210μs 16.4157μs 60.9174 KOps/s 59.2973 KOps/s $\color{#35bf28}+2.73\%$
test_plain_set_stack_nested_inplace 42.1600μs 16.3016μs 61.3435 KOps/s 60.4733 KOps/s $\color{#35bf28}+1.44\%$
test_items 32.2610μs 5.8482μs 170.9934 KOps/s 168.8200 KOps/s $\color{#35bf28}+1.29\%$
test_items_nested 0.4865ms 0.4547ms 2.1990 KOps/s 2.1983 KOps/s $\color{#35bf28}+0.03\%$
test_items_nested_locked 0.5223ms 0.4574ms 2.1864 KOps/s 2.1855 KOps/s $\color{#35bf28}+0.04\%$
test_items_nested_leaf 0.1345ms 96.6506μs 10.3465 KOps/s 10.4761 KOps/s $\color{#d91a1a}-1.24\%$
test_items_stack_nested 0.4892ms 0.4546ms 2.1995 KOps/s 2.1943 KOps/s $\color{#35bf28}+0.24\%$
test_items_stack_nested_leaf 0.1235ms 95.8758μs 10.4302 KOps/s 10.2397 KOps/s $\color{#35bf28}+1.86\%$
test_items_stack_nested_locked 0.5031ms 0.4568ms 2.1891 KOps/s 2.1795 KOps/s $\color{#35bf28}+0.44\%$
test_keys 32.2410μs 4.3939μs 227.5858 KOps/s 238.5046 KOps/s $\color{#d91a1a}-4.58\%$
test_keys_nested 0.1500ms 0.1194ms 8.3768 KOps/s 8.2991 KOps/s $\color{#35bf28}+0.94\%$
test_keys_nested_locked 2.2264ms 0.1289ms 7.7574 KOps/s 7.0254 KOps/s $\textbf{\color{#35bf28}+10.42\%}$
test_keys_nested_leaf 0.4273ms 0.1095ms 9.1364 KOps/s 8.8941 KOps/s $\color{#35bf28}+2.73\%$
test_keys_stack_nested 0.2768ms 0.1191ms 8.3970 KOps/s 8.1705 KOps/s $\color{#35bf28}+2.77\%$
test_keys_stack_nested_leaf 0.1501ms 0.1094ms 9.1376 KOps/s 8.9332 KOps/s $\color{#35bf28}+2.29\%$
test_keys_stack_nested_locked 0.1579ms 0.1279ms 7.8188 KOps/s 7.6637 KOps/s $\color{#35bf28}+2.02\%$
test_values 5.8800μs 1.0265μs 974.1928 KOps/s 970.7254 KOps/s $\color{#35bf28}+0.36\%$
test_values_nested 72.0510μs 47.5578μs 21.0271 KOps/s 20.7164 KOps/s $\color{#35bf28}+1.50\%$
test_values_nested_locked 74.7610μs 50.7543μs 19.7028 KOps/s 19.5340 KOps/s $\color{#35bf28}+0.86\%$
test_values_nested_leaf 76.6910μs 53.6596μs 18.6360 KOps/s 18.2046 KOps/s $\color{#35bf28}+2.37\%$
test_values_stack_nested 74.2000μs 47.4325μs 21.0826 KOps/s 20.5257 KOps/s $\color{#35bf28}+2.71\%$
test_values_stack_nested_leaf 82.8510μs 54.2405μs 18.4364 KOps/s 18.1921 KOps/s $\color{#35bf28}+1.34\%$
test_values_stack_nested_locked 80.8610μs 50.5523μs 19.7815 KOps/s 19.4720 KOps/s $\color{#35bf28}+1.59\%$
test_membership 5.7133μs 0.8752μs 1.1426 MOps/s 1.1638 MOps/s $\color{#d91a1a}-1.83\%$
test_membership_nested 25.6410μs 2.9243μs 341.9643 KOps/s 346.7291 KOps/s $\color{#d91a1a}-1.37\%$
test_membership_nested_leaf 34.6600μs 2.9232μs 342.0959 KOps/s 348.4674 KOps/s $\color{#d91a1a}-1.83\%$
test_membership_stacked_nested 23.7110μs 2.9144μs 343.1208 KOps/s 342.4021 KOps/s $\color{#35bf28}+0.21\%$
test_membership_stacked_nested_leaf 23.8600μs 2.9275μs 341.5926 KOps/s 344.1591 KOps/s $\color{#d91a1a}-0.75\%$
test_membership_nested_last 35.9100μs 4.3313μs 230.8770 KOps/s 229.8569 KOps/s $\color{#35bf28}+0.44\%$
test_membership_nested_leaf_last 46.7410μs 4.2972μs 232.7094 KOps/s 228.8224 KOps/s $\color{#35bf28}+1.70\%$
test_membership_stacked_nested_last 22.6300μs 4.3340μs 230.7357 KOps/s 230.7970 KOps/s $\color{#d91a1a}-0.03\%$
test_membership_stacked_nested_leaf_last 25.0800μs 4.3458μs 230.1093 KOps/s 234.0429 KOps/s $\color{#d91a1a}-1.68\%$
test_nested_getleaf 80.1210μs 21.3191μs 46.9064 KOps/s 47.5993 KOps/s $\color{#d91a1a}-1.46\%$
test_nested_get 44.6300μs 20.2624μs 49.3525 KOps/s 49.4813 KOps/s $\color{#d91a1a}-0.26\%$
test_stacked_getleaf 50.9310μs 21.6704μs 46.1458 KOps/s 46.5277 KOps/s $\color{#d91a1a}-0.82\%$
test_stacked_get 50.5510μs 20.5217μs 48.7290 KOps/s 49.6260 KOps/s $\color{#d91a1a}-1.81\%$
test_nested_getitemleaf 42.1400μs 21.8484μs 45.7700 KOps/s 46.2110 KOps/s $\color{#d91a1a}-0.95\%$
test_nested_getitem 48.7010μs 20.8533μs 47.9540 KOps/s 48.1775 KOps/s $\color{#d91a1a}-0.46\%$
test_stacked_getitemleaf 58.1710μs 21.9363μs 45.5865 KOps/s 45.9048 KOps/s $\color{#d91a1a}-0.69\%$
test_stacked_getitem 44.5700μs 20.9402μs 47.7551 KOps/s 47.9593 KOps/s $\color{#d91a1a}-0.43\%$
test_lock_nested 0.6420ms 0.4673ms 2.1400 KOps/s 2.0868 KOps/s $\color{#35bf28}+2.55\%$
test_lock_stack_nested 0.5219ms 0.4747ms 2.1067 KOps/s 2.0602 KOps/s $\color{#35bf28}+2.26\%$
test_unlock_nested 0.5365ms 0.3800ms 2.6316 KOps/s 2.6154 KOps/s $\color{#35bf28}+0.62\%$
test_unlock_stack_nested 0.4055ms 0.3812ms 2.6236 KOps/s 2.5645 KOps/s $\color{#35bf28}+2.30\%$
test_flatten_speed 0.1578ms 0.1228ms 8.1411 KOps/s 8.0898 KOps/s $\color{#35bf28}+0.63\%$
test_unflatten_speed 0.6045ms 0.5713ms 1.7503 KOps/s 1.7377 KOps/s $\color{#35bf28}+0.73\%$
test_common_ops 0.8189ms 0.6758ms 1.4797 KOps/s 1.4506 KOps/s $\color{#35bf28}+2.01\%$
test_creation 0.1224ms 2.9236μs 342.0435 KOps/s 339.9941 KOps/s $\color{#35bf28}+0.60\%$
test_creation_empty 31.1600μs 6.1693μs 162.0937 KOps/s 161.3743 KOps/s $\color{#35bf28}+0.45\%$
test_creation_nested_1 32.8500μs 10.4826μs 95.3958 KOps/s 95.2035 KOps/s $\color{#35bf28}+0.20\%$
test_creation_nested_2 37.8800μs 11.9302μs 83.8207 KOps/s 84.0326 KOps/s $\color{#d91a1a}-0.25\%$
test_creation_many_keys[10] 46.1300μs 18.3189μs 54.5885 KOps/s 54.7449 KOps/s $\color{#d91a1a}-0.29\%$
test_creation_many_keys[50] 0.1012ms 78.3686μs 12.7602 KOps/s 12.7518 KOps/s $\color{#35bf28}+0.07\%$
test_creation_many_keys[100] 0.1924ms 0.1541ms 6.4875 KOps/s 6.5927 KOps/s $\color{#d91a1a}-1.60\%$
test_creation_nested_many_keys[10] 75.0110μs 39.6116μs 25.2451 KOps/s 25.4172 KOps/s $\color{#d91a1a}-0.68\%$
test_creation_nested_many_keys[50] 0.1987ms 0.1608ms 6.2182 KOps/s 6.2920 KOps/s $\color{#d91a1a}-1.17\%$
test_clone 53.2010μs 13.1141μs 76.2539 KOps/s 73.3316 KOps/s $\color{#35bf28}+3.99\%$
test_getitem[int] 1.6699ms 14.5137μs 68.9003 KOps/s 55.2406 KOps/s $\textbf{\color{#35bf28}+24.73\%}$
test_getitem[slice_int] 0.1393ms 24.5726μs 40.6957 KOps/s 40.2487 KOps/s $\color{#35bf28}+1.11\%$
test_getitem[range] 0.1759ms 62.7166μs 15.9447 KOps/s 15.8656 KOps/s $\color{#35bf28}+0.50\%$
test_getitem[tuple] 0.1426ms 23.8964μs 41.8473 KOps/s 40.9576 KOps/s $\color{#35bf28}+2.17\%$
test_getitem[list] 0.1872ms 57.6451μs 17.3475 KOps/s 16.9944 KOps/s $\color{#35bf28}+2.08\%$
test_setitem_dim[int] 59.2000μs 26.0464μs 38.3931 KOps/s 37.6567 KOps/s $\color{#35bf28}+1.96\%$
test_setitem_dim[slice_int] 65.0010μs 43.7053μs 22.8805 KOps/s 22.8543 KOps/s $\color{#35bf28}+0.11\%$
test_setitem_dim[range] 0.1173ms 94.1894μs 10.6169 KOps/s 10.5063 KOps/s $\color{#35bf28}+1.05\%$
test_setitem_dim[tuple] 65.8010μs 40.2951μs 24.8169 KOps/s 24.4182 KOps/s $\color{#35bf28}+1.63\%$
test_setitem 71.9700μs 17.6751μs 56.5769 KOps/s 54.5869 KOps/s $\color{#35bf28}+3.65\%$
test_set 61.8410μs 16.6793μs 59.9546 KOps/s 57.5105 KOps/s $\color{#35bf28}+4.25\%$
test_set_shared 0.5664ms 0.2040ms 4.9012 KOps/s 4.8597 KOps/s $\color{#35bf28}+0.85\%$
test_update 0.2171ms 21.0206μs 47.5723 KOps/s 45.5298 KOps/s $\color{#35bf28}+4.49\%$
test_update_nested 75.3310μs 33.4101μs 29.9311 KOps/s 29.4403 KOps/s $\color{#35bf28}+1.67\%$
test_update__nested 0.4594ms 33.6717μs 29.6985 KOps/s 27.8916 KOps/s $\textbf{\color{#35bf28}+6.48\%}$
test_set_nested 60.2100μs 18.4219μs 54.2831 KOps/s 52.4798 KOps/s $\color{#35bf28}+3.44\%$
test_set_nested_new 65.7600μs 23.4964μs 42.5598 KOps/s 41.6871 KOps/s $\color{#35bf28}+2.09\%$
test_select 71.8310μs 40.4687μs 24.7104 KOps/s 24.1407 KOps/s $\color{#35bf28}+2.36\%$
test_select_nested 0.1086ms 74.9644μs 13.3397 KOps/s 13.3804 KOps/s $\color{#d91a1a}-0.30\%$
test_exclude_nested 0.1205ms 92.6055μs 10.7985 KOps/s 10.8050 KOps/s $\color{#d91a1a}-0.06\%$
test_empty[True] 0.4431ms 0.3890ms 2.5704 KOps/s 2.5464 KOps/s $\color{#35bf28}+0.94\%$
test_empty[False] 7.6952μs 1.3128μs 761.7101 KOps/s 752.0084 KOps/s $\color{#35bf28}+1.29\%$
test_to 0.1027ms 71.4021μs 14.0052 KOps/s 13.6464 KOps/s $\color{#35bf28}+2.63\%$
test_to_nonblocking 0.1008ms 64.7811μs 15.4366 KOps/s 15.4285 KOps/s $\color{#35bf28}+0.05\%$
test_unbind_speed 0.3639ms 0.3253ms 3.0743 KOps/s 3.0294 KOps/s $\color{#35bf28}+1.48\%$
test_unbind_speed_stack0 0.3715ms 0.3244ms 3.0824 KOps/s 3.0717 KOps/s $\color{#35bf28}+0.35\%$
test_unbind_speed_stack1 0.1045s 0.8292ms 1.2060 KOps/s 1.1705 KOps/s $\color{#35bf28}+3.03\%$
test_split 0.1045s 1.2716ms 786.4397 Ops/s 777.3484 Ops/s $\color{#35bf28}+1.17\%$
test_chunk 0.1040s 1.2077ms 828.0387 Ops/s 915.1694 Ops/s $\textbf{\color{#d91a1a}-9.52\%}$
test_to_cpu_blocking 19.1893ms 19.0764ms 52.4207 Ops/s 31.5390 Ops/s $\textbf{\color{#35bf28}+66.21\%}$
test_to_cpu_global_sync 11.1104ms 10.9896ms 90.9955 Ops/s 88.8616 Ops/s $\color{#35bf28}+2.40\%$
test_to_cpu_event_sync 12.2831ms 12.0066ms 83.2875 Ops/s 81.6599 Ops/s $\color{#35bf28}+1.99\%$
test_to_cpu_default 0.1162s 13.2404ms 75.5262 Ops/s 81.6999 Ops/s $\textbf{\color{#d91a1a}-7.56\%}$
test_consolidate[False-None] 4.3047ms 4.1481ms 241.0762 Ops/s 216.7138 Ops/s $\textbf{\color{#35bf28}+11.24\%}$
test_consolidate[default-None] 2.1748ms 2.0470ms 488.5201 Ops/s 483.4830 Ops/s $\color{#35bf28}+1.04\%$
test_consolidate[reduce-overhead-None] 2.1005ms 1.9634ms 509.3105 Ops/s 500.7562 Ops/s $\color{#35bf28}+1.71\%$
test_consolidate_njt[False-None] 8.6566ms 8.4623ms 118.1714 Ops/s 116.5662 Ops/s $\color{#35bf28}+1.38\%$
test_to[False-False-None] 2.1532ms 2.0467ms 488.5971 Ops/s 480.9401 Ops/s $\color{#35bf28}+1.59\%$
test_to[True-False-None] 2.1986ms 1.9457ms 513.9577 Ops/s 525.8979 Ops/s $\color{#d91a1a}-2.27\%$
test_to[within-False-None] 6.2895ms 6.1573ms 162.4094 Ops/s 161.7854 Ops/s $\color{#35bf28}+0.39\%$
test_to[True-default-None] 7.6726ms 7.5595ms 132.2847 Ops/s 130.0377 Ops/s $\color{#35bf28}+1.73\%$
test_to_njt[False-False-None] 8.7126ms 8.5342ms 117.1761 Ops/s 115.5047 Ops/s $\color{#35bf28}+1.45\%$
test_to_njt[True-False-None] 7.1689ms 6.9997ms 142.8624 Ops/s 141.4719 Ops/s $\color{#35bf28}+0.98\%$
test_to_njt[within-False-None] 15.9317ms 15.7642ms 63.4347 Ops/s 63.1094 Ops/s $\color{#35bf28}+0.52\%$
test_creation[device0] 0.4469ms 0.1174ms 8.5172 KOps/s 8.1877 KOps/s $\color{#35bf28}+4.02\%$
test_creation_from_tensor 0.4568ms 0.1139ms 8.7775 KOps/s 8.4509 KOps/s $\color{#35bf28}+3.86\%$
test_add_one[memmap_tensor0] 0.4256ms 6.4851μs 154.1997 KOps/s 149.3105 KOps/s $\color{#35bf28}+3.27\%$
test_contiguous[memmap_tensor0] 31.4410μs 0.6686μs 1.4956 MOps/s 2.2275 MOps/s $\textbf{\color{#d91a1a}-32.86\%}$
test_stack[memmap_tensor0] 57.1200μs 4.6354μs 215.7306 KOps/s 209.0810 KOps/s $\color{#35bf28}+3.18\%$
test_memmaptd_index 1.0860ms 0.2663ms 3.7559 KOps/s 3.7287 KOps/s $\color{#35bf28}+0.73\%$
test_memmaptd_index_astensor 0.5118ms 0.3581ms 2.7926 KOps/s 2.7464 KOps/s $\color{#35bf28}+1.68\%$
test_memmaptd_index_op 0.8725ms 0.6044ms 1.6545 KOps/s 1.6309 KOps/s $\color{#35bf28}+1.45\%$
test_serialize_model 0.1387s 0.1364s 7.3322 Ops/s 7.3484 Ops/s $\color{#d91a1a}-0.22\%$
test_serialize_model_pickle 1.3472s 1.1926s 0.8385 Ops/s 0.8260 Ops/s $\color{#35bf28}+1.51\%$
test_serialize_weights 0.1408s 0.1347s 7.4266 Ops/s 7.3939 Ops/s $\color{#35bf28}+0.44\%$
test_serialize_weights_returnearly 0.3757s 93.9065ms 10.6489 Ops/s 14.6216 Ops/s $\textbf{\color{#d91a1a}-27.17\%}$
test_serialize_weights_pickle 1.3657s 1.1883s 0.8415 Ops/s 0.8245 Ops/s $\color{#35bf28}+2.06\%$
test_reshape_pytree 0.2133ms 33.7608μs 29.6202 KOps/s 29.5971 KOps/s $\color{#35bf28}+0.08\%$
test_reshape_td 73.1110μs 44.2177μs 22.6154 KOps/s 22.2993 KOps/s $\color{#35bf28}+1.42\%$
test_view_pytree 0.2526ms 33.4013μs 29.9389 KOps/s 29.6836 KOps/s $\color{#35bf28}+0.86\%$
test_view_td 82.1710μs 51.8336μs 19.2925 KOps/s 18.7960 KOps/s $\color{#35bf28}+2.64\%$
test_unbind_pytree 0.2366ms 37.1554μs 26.9140 KOps/s 26.4851 KOps/s $\color{#35bf28}+1.62\%$
test_unbind_td 0.1174ms 49.6290μs 20.1495 KOps/s 20.3667 KOps/s $\color{#d91a1a}-1.07\%$
test_split_pytree 0.2556ms 42.8721μs 23.3252 KOps/s 23.0854 KOps/s $\color{#35bf28}+1.04\%$
test_split_td 0.1343ms 64.3527μs 15.5394 KOps/s 15.4168 KOps/s $\color{#35bf28}+0.79\%$
test_add_pytree 0.2311ms 42.5488μs 23.5024 KOps/s 22.8257 KOps/s $\color{#35bf28}+2.96\%$
test_add_td 92.3620μs 53.0258μs 18.8587 KOps/s 17.9073 KOps/s $\textbf{\color{#35bf28}+5.31\%}$
test_compile_add_one_nested[tensordict-compile] 0.2079ms 0.1385ms 7.2198 KOps/s 6.9448 KOps/s $\color{#35bf28}+3.96\%$
test_compile_add_one_nested[tensordict-eager] 0.3140ms 0.1940ms 5.1551 KOps/s 5.2684 KOps/s $\color{#d91a1a}-2.15\%$
test_compile_add_one_nested[pytree-compile] 0.1763ms 0.1081ms 9.2511 KOps/s 8.9448 KOps/s $\color{#35bf28}+3.42\%$
test_compile_add_one_nested[pytree-eager] 0.4766ms 0.1842ms 5.4295 KOps/s 5.4278 KOps/s $\color{#35bf28}+0.03\%$
test_compile_copy_nested[tensordict-compile] 70.6510μs 30.6021μs 32.6775 KOps/s 33.0887 KOps/s $\color{#d91a1a}-1.24\%$
test_compile_copy_nested[tensordict-eager] 93.6610μs 52.6545μs 18.9917 KOps/s 19.0549 KOps/s $\color{#d91a1a}-0.33\%$
test_compile_copy_nested[pytree-compile] 44.4410μs 9.9010μs 101.0000 KOps/s 100.2981 KOps/s $\color{#35bf28}+0.70\%$
test_compile_copy_nested[pytree-eager] 0.4279ms 70.4041μs 14.2037 KOps/s 14.0334 KOps/s $\color{#35bf28}+1.21\%$
test_compile_add_one_flat[tensordict-compile] 0.2320ms 0.1772ms 5.6434 KOps/s 5.1738 KOps/s $\textbf{\color{#35bf28}+9.08\%}$
test_compile_add_one_flat[tensordict-eager] 0.3756ms 0.2529ms 3.9542 KOps/s 3.8745 KOps/s $\color{#35bf28}+2.06\%$
test_compile_add_one_flat[tensorclass-compile] 0.1653ms 0.1159ms 8.6290 KOps/s 8.0431 KOps/s $\textbf{\color{#35bf28}+7.28\%}$
test_compile_add_one_flat[tensorclass-eager] 0.1258ms 68.8639μs 14.5214 KOps/s 14.3198 KOps/s $\color{#35bf28}+1.41\%$
test_compile_add_one_flat[pytree-compile] 0.2165ms 0.1588ms 6.2985 KOps/s 6.0294 KOps/s $\color{#35bf28}+4.46\%$
test_compile_add_one_flat[pytree-eager] 0.8526ms 0.5262ms 1.9005 KOps/s 1.8270 KOps/s $\color{#35bf28}+4.03\%$
test_compile_add_self_flat[tensordict-eager] 0.3547ms 0.3058ms 3.2703 KOps/s 3.1821 KOps/s $\color{#35bf28}+2.77\%$
test_compile_add_self_flat[tensordict-compile] 0.3425ms 0.1859ms 5.3786 KOps/s 5.0678 KOps/s $\textbf{\color{#35bf28}+6.13\%}$
test_compile_add_self_flat[tensorclass-eager] 0.2794ms 85.4661μs 11.7005 KOps/s 11.8334 KOps/s $\color{#d91a1a}-1.12\%$
test_compile_add_self_flat[tensorclass-compile] 0.1600ms 0.1175ms 8.5087 KOps/s 7.7757 KOps/s $\textbf{\color{#35bf28}+9.43\%}$
test_compile_add_self_flat[pytree-eager] 0.6598ms 0.4366ms 2.2902 KOps/s 2.2264 KOps/s $\color{#35bf28}+2.87\%$
test_compile_add_self_flat[pytree-compile] 0.2154ms 0.1583ms 6.3168 KOps/s 6.0378 KOps/s $\color{#35bf28}+4.62\%$
test_compile_copy_flat[tensordict-compile] 91.4010μs 24.4982μs 40.8193 KOps/s 37.0589 KOps/s $\textbf{\color{#35bf28}+10.15\%}$
test_compile_copy_flat[tensordict-eager] 82.7710μs 41.6308μs 24.0207 KOps/s 23.9978 KOps/s $\color{#35bf28}+0.10\%$
test_compile_copy_flat[pytree-compile] 45.1910μs 10.8997μs 91.7457 KOps/s 91.2452 KOps/s $\color{#35bf28}+0.55\%$
test_compile_copy_flat[pytree-eager] 0.3862ms 52.2609μs 19.1348 KOps/s 19.2242 KOps/s $\color{#d91a1a}-0.47\%$
test_compile_assign_and_add[tensordict-compile] 2.0334ms 0.1736ms 5.7598 KOps/s 5.3468 KOps/s $\textbf{\color{#35bf28}+7.72\%}$
test_compile_assign_and_add[tensordict-eager] 3.5163ms 3.2988ms 303.1371 Ops/s 297.8432 Ops/s $\color{#35bf28}+1.78\%$
test_compile_assign_and_add[pytree-compile] 2.0038ms 0.1619ms 6.1754 KOps/s 5.8403 KOps/s $\textbf{\color{#35bf28}+5.74\%}$
test_compile_assign_and_add[pytree-eager] 2.9075ms 2.7651ms 361.6549 Ops/s 353.4527 Ops/s $\color{#35bf28}+2.32\%$
test_compile_indexing[tensor-tensordict-compile] 0.1685ms 0.1117ms 8.9561 KOps/s 8.7543 KOps/s $\color{#35bf28}+2.30\%$
test_compile_indexing[tensor-tensordict-eager] 0.3238ms 76.5610μs 13.0615 KOps/s 13.3347 KOps/s $\color{#d91a1a}-2.05\%$
test_compile_indexing[tensor-tensorclass-compile] 0.1431ms 96.3250μs 10.3815 KOps/s 10.0575 KOps/s $\color{#35bf28}+3.22\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2640ms 47.1551μs 21.2066 KOps/s 21.4620 KOps/s $\color{#d91a1a}-1.19\%$
test_compile_indexing[tensor-pytree-compile] 0.1449ms 96.6449μs 10.3472 KOps/s 10.0948 KOps/s $\color{#35bf28}+2.50\%$
test_compile_indexing[tensor-pytree-eager] 0.2629ms 46.9174μs 21.3141 KOps/s 16.2020 KOps/s $\textbf{\color{#35bf28}+31.55\%}$
test_compile_indexing[slice-tensordict-compile] 0.1033ms 59.0415μs 16.9372 KOps/s 16.7354 KOps/s $\color{#35bf28}+1.21\%$
test_compile_indexing[slice-tensordict-eager] 0.2487ms 28.3937μs 35.2190 KOps/s 35.7888 KOps/s $\color{#d91a1a}-1.59\%$
test_compile_indexing[slice-tensorclass-compile] 91.7010μs 45.7263μs 21.8693 KOps/s 21.6439 KOps/s $\color{#35bf28}+1.04\%$
test_compile_indexing[slice-tensorclass-eager] 0.2546ms 22.7397μs 43.9760 KOps/s 43.6516 KOps/s $\color{#35bf28}+0.74\%$
test_compile_indexing[slice-pytree-compile] 82.1610μs 45.4322μs 22.0108 KOps/s 21.1889 KOps/s $\color{#35bf28}+3.88\%$
test_compile_indexing[slice-pytree-eager] 0.2896ms 23.3441μs 42.8373 KOps/s 43.7857 KOps/s $\color{#d91a1a}-2.17\%$
test_compile_indexing[int-tensordict-compile] 0.1471ms 56.7202μs 17.6304 KOps/s 16.1718 KOps/s $\textbf{\color{#35bf28}+9.02\%}$
test_compile_indexing[int-tensordict-eager] 0.3259ms 27.3837μs 36.5181 KOps/s 36.1386 KOps/s $\color{#35bf28}+1.05\%$
test_compile_indexing[int-tensorclass-compile] 99.6110μs 47.3749μs 21.1082 KOps/s 21.2361 KOps/s $\color{#d91a1a}-0.60\%$
test_compile_indexing[int-tensorclass-eager] 0.2647ms 22.6843μs 44.0834 KOps/s 43.6273 KOps/s $\color{#35bf28}+1.05\%$
test_compile_indexing[int-pytree-compile] 0.1002ms 45.4230μs 22.0153 KOps/s 20.6393 KOps/s $\textbf{\color{#35bf28}+6.67\%}$
test_compile_indexing[int-pytree-eager] 0.2634ms 23.0848μs 43.3185 KOps/s 43.7852 KOps/s $\color{#d91a1a}-1.07\%$
test_mod_add[eager] 98.9420μs 51.6466μs 19.3624 KOps/s 19.6854 KOps/s $\color{#d91a1a}-1.64\%$
test_mod_add[compile] 0.1537ms 0.1085ms 9.2136 KOps/s 9.2752 KOps/s $\color{#d91a1a}-0.66\%$
test_mod_add[compile-overhead] 0.2367ms 0.1495ms 6.6901 KOps/s 6.5222 KOps/s $\color{#35bf28}+2.57\%$
test_mod_wrap[eager] 0.3679ms 0.2946ms 3.3950 KOps/s 3.3953 KOps/s $-0.01\%$
test_mod_wrap[compile] 0.4141ms 0.3462ms 2.8882 KOps/s 2.8064 KOps/s $\color{#35bf28}+2.91\%$
test_mod_wrap[compile-overhead] 7.3676ms 4.0520ms 246.7932 Ops/s 248.1377 Ops/s $\color{#d91a1a}-0.54\%$
test_mod_wrap_and_backward[eager] 1.5940ms 1.4932ms 669.7190 Ops/s 620.9907 Ops/s $\textbf{\color{#35bf28}+7.85\%}$
test_mod_wrap_and_backward[compile] 1.5177ms 1.4387ms 695.0627 Ops/s 635.1017 Ops/s $\textbf{\color{#35bf28}+9.44\%}$
test_mod_wrap_and_backward[compile-overhead] 1.2597ms 0.8895ms 1.1243 KOps/s 987.6368 Ops/s $\textbf{\color{#35bf28}+13.83\%}$
test_seq_add[eager] 0.2281ms 0.1561ms 6.4044 KOps/s 6.1516 KOps/s $\color{#35bf28}+4.11\%$
test_seq_add[compile] 0.1727ms 0.1178ms 8.4875 KOps/s 8.0126 KOps/s $\textbf{\color{#35bf28}+5.93\%}$
test_seq_add[compile-overhead] 0.3629ms 0.1597ms 6.2635 KOps/s 6.1270 KOps/s $\color{#35bf28}+2.23\%$
test_seq_wrap[eager] 0.6208ms 0.5405ms 1.8502 KOps/s 1.8478 KOps/s $\color{#35bf28}+0.13\%$
test_seq_wrap[compile] 0.4179ms 0.3646ms 2.7424 KOps/s 2.5624 KOps/s $\textbf{\color{#35bf28}+7.03\%}$
test_seq_wrap[compile-overhead] 0.3359ms 0.2644ms 3.7825 KOps/s 3.6393 KOps/s $\color{#35bf28}+3.94\%$
test_func_call_runtime[False-eager] 0.9160ms 0.8450ms 1.1835 KOps/s 1.1294 KOps/s $\color{#35bf28}+4.78\%$
test_func_call_runtime[False-compile] 2.1767ms 0.9120ms 1.0964 KOps/s 1.0809 KOps/s $\color{#35bf28}+1.44\%$
test_func_call_runtime[False-compile-overhead] 0.5107ms 0.4608ms 2.1701 KOps/s 2.1281 KOps/s $\color{#35bf28}+1.97\%$
test_func_call_runtime[True-eager] 1.1565ms 1.0824ms 923.8512 Ops/s 891.1624 Ops/s $\color{#35bf28}+3.67\%$
test_func_call_runtime[True-compile] 0.9975ms 0.9184ms 1.0889 KOps/s 1.0547 KOps/s $\color{#35bf28}+3.24\%$
test_func_call_runtime[True-compile-overhead] 0.6077ms 0.4750ms 2.1052 KOps/s 2.0573 KOps/s $\color{#35bf28}+2.32\%$
test_func_call_cm_runtime[False-eager] 1.0703ms 0.8825ms 1.1332 KOps/s 1.1790 KOps/s $\color{#d91a1a}-3.88\%$
test_func_call_cm_runtime[False-compile] 1.1222ms 0.9104ms 1.0984 KOps/s 1.0729 KOps/s $\color{#35bf28}+2.38\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5700ms 0.4639ms 2.1555 KOps/s 2.1333 KOps/s $\color{#35bf28}+1.04\%$
test_func_call_cm_runtime[True-eager] 1.4812ms 1.2323ms 811.5149 Ops/s 798.6816 Ops/s $\color{#35bf28}+1.61\%$
test_func_call_cm_runtime[True-compile] 1.0367ms 0.9546ms 1.0476 KOps/s 1.0230 KOps/s $\color{#35bf28}+2.40\%$
test_func_call_cm_runtime[True-compile-overhead] 0.5555ms 0.5062ms 1.9754 KOps/s 1.9401 KOps/s $\color{#35bf28}+1.82\%$
test_vmap_func_call_cm_runtime[eager] 2.8455ms 2.3574ms 424.2036 Ops/s 417.2925 Ops/s $\color{#35bf28}+1.66\%$
test_vmap_func_call_cm_runtime[compile] 1.0342ms 0.9710ms 1.0299 KOps/s 1.0091 KOps/s $\color{#35bf28}+2.06\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.5764ms 0.5120ms 1.9531 KOps/s 1.9127 KOps/s $\color{#35bf28}+2.11\%$
test_distributed 2.5648ms 0.1585ms 6.3110 KOps/s 6.4635 KOps/s $\color{#d91a1a}-2.36\%$
test_tdmodule 0.5367ms 28.4505μs 35.1488 KOps/s 35.5581 KOps/s $\color{#d91a1a}-1.15\%$
test_tdmodule_dispatch 76.0910μs 45.5119μs 21.9723 KOps/s 21.5537 KOps/s $\color{#35bf28}+1.94\%$
test_tdseq 46.6700μs 26.9832μs 37.0601 KOps/s 36.6511 KOps/s $\color{#35bf28}+1.12\%$
test_tdseq_dispatch 80.2710μs 46.9970μs 21.2780 KOps/s 20.7663 KOps/s $\color{#35bf28}+2.46\%$
test_instantiation_functorch 2.1531ms 2.0599ms 485.4651 Ops/s 478.9384 Ops/s $\color{#35bf28}+1.36\%$
test_exec_functorch 0.2256ms 0.1810ms 5.5262 KOps/s 5.5645 KOps/s $\color{#d91a1a}-0.69\%$
test_exec_functional_call 0.1974ms 0.1609ms 6.2133 KOps/s 6.1571 KOps/s $\color{#35bf28}+0.91\%$
test_exec_td_decorator 0.4534ms 0.2369ms 4.2213 KOps/s 4.1880 KOps/s $\color{#35bf28}+0.80\%$
test_vmap_mlp_speed_decorator[True-True] 1.0071ms 0.8186ms 1.2217 KOps/s 1.2011 KOps/s $\color{#35bf28}+1.71\%$
test_vmap_mlp_speed_decorator[True-False] 1.0184ms 0.8179ms 1.2226 KOps/s 1.2046 KOps/s $\color{#35bf28}+1.49\%$
test_vmap_mlp_speed_decorator[False-True] 0.8677ms 0.7117ms 1.4051 KOps/s 1.3878 KOps/s $\color{#35bf28}+1.25\%$
test_vmap_mlp_speed_decorator[False-False] 0.9091ms 0.7112ms 1.4061 KOps/s 1.3935 KOps/s $\color{#35bf28}+0.90\%$
test_vmap_transformer_speed_decorator[True-True] 21.5032ms 20.4717ms 48.8480 Ops/s 48.3624 Ops/s $\color{#35bf28}+1.00\%$
test_vmap_transformer_speed_decorator[True-False] 21.0272ms 20.4433ms 48.9158 Ops/s 48.3853 Ops/s $\color{#35bf28}+1.10\%$
test_vmap_transformer_speed_decorator[False-True] 20.3845ms 20.2487ms 49.3860 Ops/s 48.7421 Ops/s $\color{#35bf28}+1.32\%$
test_vmap_transformer_speed_decorator[False-False] 20.9857ms 20.2961ms 49.2705 Ops/s 48.7840 Ops/s $\color{#35bf28}+1.00\%$
test_to_module_speed[True] 1.5760ms 1.4899ms 671.1729 Ops/s 680.1013 Ops/s $\color{#d91a1a}-1.31\%$
test_to_module_speed[False] 1.5545ms 1.4627ms 683.6779 Ops/s 689.3720 Ops/s $\color{#d91a1a}-0.83\%$
test_tc_init 80.8410μs 45.0203μs 22.2122 KOps/s 21.9499 KOps/s $\color{#35bf28}+1.20\%$
test_tc_init_tensor_only 35.5400μs 9.8438μs 101.5872 KOps/s 101.0610 KOps/s $\color{#35bf28}+0.52\%$
test_tc_init_nested 0.1338ms 89.8636μs 11.1280 KOps/s 11.0633 KOps/s $\color{#35bf28}+0.58\%$
test_tc_init_many_fields 43.8910μs 16.7990μs 59.5273 KOps/s 60.7700 KOps/s $\color{#d91a1a}-2.04\%$
test_tc_first_layer_tensor 18.1000μs 1.8702μs 534.6992 KOps/s 544.1977 KOps/s $\color{#d91a1a}-1.75\%$
test_tc_first_layer_tensor_only 8.9287μs 0.7679μs 1.3023 MOps/s 1.3329 MOps/s $\color{#d91a1a}-2.29\%$
test_tc_first_layer_tensor_set 24.1400μs 3.9549μs 252.8489 KOps/s 252.6352 KOps/s $\color{#35bf28}+0.08\%$
test_tc_first_layer_tensor_only_set 29.3800μs 3.1133μs 321.2027 KOps/s 318.1045 KOps/s $\color{#35bf28}+0.97\%$
test_tc_first_layer_nontensor 25.9300μs 6.2719μs 159.4413 KOps/s 155.2762 KOps/s $\color{#35bf28}+2.68\%$
test_tc_second_layer_tensor 24.7510μs 4.4221μs 226.1350 KOps/s 228.5019 KOps/s $\color{#d91a1a}-1.04\%$
test_tc_second_layer_nontensor 36.9600μs 8.6238μs 115.9580 KOps/s 109.3909 KOps/s $\textbf{\color{#35bf28}+6.00\%}$
test_unbind 0.2468s 14.0324ms 71.2635 Ops/s 55.2049 Ops/s $\textbf{\color{#35bf28}+29.09\%}$
test_full_like 9.3068ms 7.3427ms 136.1905 Ops/s 228.7819 Ops/s $\textbf{\color{#d91a1a}-40.47\%}$
test_zeros_like 4.9627ms 4.3606ms 229.3280 Ops/s 229.2154 Ops/s $\color{#35bf28}+0.05\%$
test_ones_like 4.4964ms 4.3612ms 229.2930 Ops/s 229.0003 Ops/s $\color{#35bf28}+0.13\%$
test_clone 6.5903ms 6.4022ms 156.1964 Ops/s 155.8008 Ops/s $\color{#35bf28}+0.25\%$
test_squeeze 0.1390ms 14.4407μs 69.2486 KOps/s 69.0852 KOps/s $\color{#35bf28}+0.24\%$
test_unsqueeze 0.2523ms 0.1111ms 8.9993 KOps/s 9.0204 KOps/s $\color{#d91a1a}-0.23\%$
test_split 0.2312ms 0.1845ms 5.4196 KOps/s 5.4081 KOps/s $\color{#35bf28}+0.21\%$
test_permute 0.2623ms 0.2019ms 4.9540 KOps/s 4.8536 KOps/s $\color{#35bf28}+2.07\%$
test_stack 51.5429ms 51.2394ms 19.5162 Ops/s 19.5229 Ops/s $\color{#d91a1a}-0.03\%$
test_cat 51.9415ms 51.2992ms 19.4935 Ops/s 19.5752 Ops/s $\color{#d91a1a}-0.42\%$
test_sequential_tensordict 0.2698ms 0.2140ms 4.6731 KOps/s 4.3964 KOps/s $\textbf{\color{#35bf28}+6.30\%}$
test_sequential_graph_module 0.5481ms 0.1189ms 8.4083 KOps/s 7.8072 KOps/s $\textbf{\color{#35bf28}+7.70\%}$
test_nested_tensordict 0.3400ms 0.2821ms 3.5442 KOps/s 3.4465 KOps/s $\color{#35bf28}+2.84\%$
test_nested_graph_module 0.5237ms 0.1274ms 7.8487 KOps/s 7.3580 KOps/s $\textbf{\color{#35bf28}+6.67\%}$

[ghstack-poisoned]
vmoens added a commit that referenced this pull request Feb 19, 2026
Send the fully qualified tensorclass type string (e.g.
``tensordict.testing.MyDistData``) as metadata during init_remote and
broadcast.  On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.

- Add ``_tensorclass_type`` private param to ``init_remote`` / ``broadcast``
  in TensorDictBase, injected into consolidation metadata.
- Add ``_resolve_tensorclass_type`` helper to import a class from its
  dotted path.
- Override ``init_remote``, ``broadcast``, ``from_remote_init`` on
  tensorclass to inject / recover the type automatically.
- Create ``tensordict/testing.py`` with ``MyDistData`` so both sides of
  a multiprocessing test can import the same class.
- Update all tensorclass distributed tests to assert ``type(result) is
  MyDistData``.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: d932b1f
Pull-Request: #1603
@github-actions

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [Distributed].

Current title: [Distributed] Preserve tensorclass type across the wire

Supported Prefixes

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[BugFix] BugFix [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Perf [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] Setup [Setup] Update build configuration

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

vmoens added a commit that referenced this pull request Feb 19, 2026
Send the fully qualified tensorclass type string (e.g.
``tensordict.testing.MyDistData``) as metadata during init_remote and
broadcast.  On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.

- Add ``_tensorclass_type`` private param to ``init_remote`` / ``broadcast``
  in TensorDictBase, injected into consolidation metadata.
- Add ``_resolve_tensorclass_type`` helper to import a class from its
  dotted path.
- Override ``init_remote``, ``broadcast``, ``from_remote_init`` on
  tensorclass to inject / recover the type automatically.
- Create ``tensordict/testing.py`` with ``MyDistData`` so both sides of
  a multiprocessing test can import the same class.
- Update all tensorclass distributed tests to assert ``type(result) is
  MyDistData``.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: d932b1f
Pull-Request: #1603
@vmoens
vmoens merged commit f09087d into gh/vmoens/81/base Feb 19, 2026
36 of 42 checks passed
@vmoens
vmoens deleted the gh/vmoens/81/head branch February 19, 2026 14:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Distributed tensorclass Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant