[Distributed] Preserve tensorclass type across the wire - #1603
Conversation
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
PR Title Label ErrorUnknown or invalid prefix Current title: Supported PrefixesYour PR title must start with exactly one of these prefixes:
Note: Common variations like singular/plural are supported (e.g., |
PR Title Label ErrorUnknown or invalid prefix Current title: Supported PrefixesYour PR title must start with exactly one of these prefixes:
Note: Common variations like singular/plural are supported (e.g., |
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
PR Title Label ErrorUnknown or invalid prefix Current title: Supported PrefixesYour PR title must start with exactly one of these prefixes:
Note: Common variations like singular/plural are supported (e.g., |
|
| 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 | |
| test_plain_set_stack_nested | 40.1310μs | 15.1518μs | 65.9987 KOps/s | 66.7656 KOps/s | |
| test_plain_set_nested_inplace | 45.9310μs | 16.3514μs | 61.1570 KOps/s | 61.3108 KOps/s | |
| test_plain_set_stack_nested_inplace | 49.7210μs | 16.2885μs | 61.3929 KOps/s | 60.8887 KOps/s | |
| test_items | 58.5610μs | 6.0112μs | 166.3559 KOps/s | 174.5380 KOps/s | |
| test_items_nested | 0.5256ms | 0.4579ms | 2.1839 KOps/s | 2.1768 KOps/s | |
| test_items_nested_locked | 0.6330ms | 0.4551ms | 2.1975 KOps/s | 2.1464 KOps/s | |
| test_items_nested_leaf | 0.1370ms | 97.0545μs | 10.3035 KOps/s | 10.2807 KOps/s | |
| test_items_stack_nested | 0.5135ms | 0.4526ms | 2.2092 KOps/s | 2.1923 KOps/s | |
| test_items_stack_nested_leaf | 0.1311ms | 95.8592μs | 10.4320 KOps/s | 10.2575 KOps/s | |
| test_items_stack_nested_locked | 0.5177ms | 0.4587ms | 2.1799 KOps/s | 2.1563 KOps/s | |
| test_keys | 28.6710μs | 4.1959μs | 238.3253 KOps/s | 237.3137 KOps/s | |
| test_keys_nested | 0.1700ms | 0.1207ms | 8.2851 KOps/s | 8.3127 KOps/s | |
| test_keys_nested_locked | 2.2105ms | 0.1299ms | 7.6986 KOps/s | 7.0529 KOps/s | |
| test_keys_nested_leaf | 0.1699ms | 0.1109ms | 9.0194 KOps/s | 8.9526 KOps/s | |
| test_keys_stack_nested | 0.1733ms | 0.1210ms | 8.2660 KOps/s | 8.2949 KOps/s | |
| test_keys_stack_nested_leaf | 0.1471ms | 0.1120ms | 8.9319 KOps/s | 9.0503 KOps/s | |
| test_keys_stack_nested_locked | 0.1646ms | 0.1285ms | 7.7806 KOps/s | 7.7533 KOps/s | |
| test_values | 6.1260μs | 1.0330μs | 968.0366 KOps/s | 986.4791 KOps/s | |
| test_values_nested | 77.4320μs | 47.5635μs | 21.0245 KOps/s | 21.0950 KOps/s | |
| test_values_nested_locked | 84.8310μs | 50.9576μs | 19.6241 KOps/s | 19.6824 KOps/s | |
| test_values_nested_leaf | 86.7620μs | 53.9908μs | 18.5217 KOps/s | 18.4441 KOps/s | |
| test_values_stack_nested | 71.6810μs | 47.5573μs | 21.0273 KOps/s | 20.9027 KOps/s | |
| test_values_stack_nested_leaf | 93.7020μs | 54.6911μs | 18.2845 KOps/s | 18.4162 KOps/s | |
| test_values_stack_nested_locked | 1.7498ms | 50.9647μs | 19.6214 KOps/s | 19.5958 KOps/s | |
| test_membership | 5.0250μs | 0.8481μs | 1.1791 MOps/s | 1.1889 MOps/s | |
| test_membership_nested | 0.6503ms | 2.8948μs | 345.4456 KOps/s | 342.9419 KOps/s | |
| test_membership_nested_leaf | 27.9000μs | 2.9298μs | 341.3186 KOps/s | 354.8948 KOps/s | |
| test_membership_stacked_nested | 26.2610μs | 2.9147μs | 343.0884 KOps/s | 343.6985 KOps/s | |
| test_membership_stacked_nested_leaf | 26.1110μs | 2.9528μs | 338.6633 KOps/s | 343.8696 KOps/s | |
| test_membership_nested_last | 40.0400μs | 4.3488μs | 229.9477 KOps/s | 228.6212 KOps/s | |
| test_membership_nested_leaf_last | 38.4010μs | 4.3206μs | 231.4512 KOps/s | 229.1702 KOps/s | |
| test_membership_stacked_nested_last | 38.1400μs | 4.3789μs | 228.3695 KOps/s | 232.2286 KOps/s | |
| test_membership_stacked_nested_leaf_last | 61.2110μs | 4.3363μs | 230.6127 KOps/s | 232.9536 KOps/s | |
| test_nested_getleaf | 50.0510μs | 21.7073μs | 46.0675 KOps/s | 47.9274 KOps/s | |
| test_nested_get | 73.5210μs | 20.5616μs | 48.6344 KOps/s | 50.5474 KOps/s | |
| test_stacked_getleaf | 55.3510μs | 21.6017μs | 46.2927 KOps/s | 47.8468 KOps/s | |
| test_stacked_get | 46.7210μs | 20.4593μs | 48.8775 KOps/s | 50.3856 KOps/s | |
| test_nested_getitemleaf | 53.6810μs | 22.3751μs | 44.6925 KOps/s | 47.3050 KOps/s | |
| test_nested_getitem | 48.7200μs | 21.1130μs | 47.3641 KOps/s | 49.6239 KOps/s | |
| test_stacked_getitemleaf | 0.1054ms | 21.9232μs | 45.6137 KOps/s | 47.0451 KOps/s | |
| test_stacked_getitem | 47.1710μs | 21.0956μs | 47.4033 KOps/s | 49.7503 KOps/s | |
| test_lock_nested | 0.7240ms | 0.4716ms | 2.1203 KOps/s | 2.0805 KOps/s | |
| test_lock_stack_nested | 0.5342ms | 0.4780ms | 2.0918 KOps/s | 2.0791 KOps/s | |
| test_unlock_nested | 0.6034ms | 0.3805ms | 2.6279 KOps/s | 2.6185 KOps/s | |
| test_unlock_stack_nested | 0.4358ms | 0.3841ms | 2.6032 KOps/s | 2.5808 KOps/s | |
| test_flatten_speed | 0.1884ms | 0.1238ms | 8.0781 KOps/s | 8.1550 KOps/s | |
| test_unflatten_speed | 0.6158ms | 0.5687ms | 1.7585 KOps/s | 1.7436 KOps/s | |
| test_common_ops | 0.8269ms | 0.6910ms | 1.4472 KOps/s | 1.4474 KOps/s | |
| test_creation | 67.9110μs | 2.9366μs | 340.5277 KOps/s | 340.5763 KOps/s | |
| test_creation_empty | 37.3600μs | 6.1475μs | 162.6690 KOps/s | 159.3918 KOps/s | |
| test_creation_nested_1 | 39.6210μs | 10.4552μs | 95.6458 KOps/s | 95.2344 KOps/s | |
| test_creation_nested_2 | 43.5910μs | 11.8577μs | 84.3332 KOps/s | 84.0492 KOps/s | |
| test_creation_many_keys[10] | 46.6110μs | 18.2447μs | 54.8105 KOps/s | 54.0859 KOps/s | |
| test_creation_many_keys[50] | 0.1081ms | 77.9439μs | 12.8297 KOps/s | 12.7910 KOps/s | |
| test_creation_many_keys[100] | 0.1963ms | 0.1525ms | 6.5570 KOps/s | 6.5386 KOps/s | |
| test_creation_nested_many_keys[10] | 71.9810μs | 39.1418μs | 25.5482 KOps/s | 25.5653 KOps/s | |
| test_creation_nested_many_keys[50] | 0.2011ms | 0.1609ms | 6.2154 KOps/s | 6.2276 KOps/s | |
| test_clone | 45.0510μs | 13.5765μs | 73.6565 KOps/s | 73.6061 KOps/s | |
| test_getitem[int] | 1.7113ms | 14.6250μs | 68.3759 KOps/s | 57.1400 KOps/s | |
| test_getitem[slice_int] | 0.1444ms | 26.6540μs | 37.5178 KOps/s | 40.6287 KOps/s | |
| test_getitem[range] | 0.1748ms | 62.8313μs | 15.9156 KOps/s | 15.3636 KOps/s | |
| test_getitem[tuple] | 0.1488ms | 23.9813μs | 41.6991 KOps/s | 42.2646 KOps/s | |
| test_getitem[list] | 0.1890ms | 61.0387μs | 16.3831 KOps/s | 16.5656 KOps/s | |
| test_setitem_dim[int] | 48.4210μs | 27.8113μs | 35.9567 KOps/s | 36.2696 KOps/s | |
| test_setitem_dim[slice_int] | 69.6610μs | 44.9166μs | 22.2635 KOps/s | 22.5365 KOps/s | |
| test_setitem_dim[range] | 0.1276ms | 98.3358μs | 10.1692 KOps/s | 10.0726 KOps/s | |
| test_setitem_dim[tuple] | 65.1310μs | 42.7804μs | 23.3752 KOps/s | 23.3336 KOps/s | |
| test_setitem | 71.6220μs | 17.8289μs | 56.0887 KOps/s | 55.0444 KOps/s | |
| test_set | 66.9520μs | 17.1116μs | 58.4399 KOps/s | 57.9852 KOps/s | |
| test_set_shared | 0.4865ms | 0.2045ms | 4.8903 KOps/s | 4.7073 KOps/s | |
| test_update | 0.3234ms | 21.5215μs | 46.4651 KOps/s | 43.3864 KOps/s | |
| test_update_nested | 64.8610μs | 33.4952μs | 29.8550 KOps/s | 28.5648 KOps/s | |
| test_update__nested | 0.4641ms | 34.4249μs | 29.0487 KOps/s | 28.4719 KOps/s | |
| test_set_nested | 60.4620μs | 19.2294μs | 52.0038 KOps/s | 51.2241 KOps/s | |
| test_set_nested_new | 63.4920μs | 23.7763μs | 42.0587 KOps/s | 38.6229 KOps/s | |
| test_select | 85.0820μs | 41.5522μs | 24.0661 KOps/s | 23.4400 KOps/s | |
| test_select_nested | 0.1087ms | 74.8673μs | 13.3570 KOps/s | 13.5005 KOps/s | |
| test_exclude_nested | 0.1274ms | 91.9732μs | 10.8727 KOps/s | 10.9001 KOps/s | |
| test_empty[True] | 0.4471ms | 0.3860ms | 2.5904 KOps/s | 2.5448 KOps/s | |
| test_empty[False] | 9.9277μs | 1.3281μs | 752.9502 KOps/s | 753.1475 KOps/s | |
| test_to | 0.1031ms | 71.0315μs | 14.0783 KOps/s | 12.6412 KOps/s | |
| test_to_nonblocking | 0.1075ms | 66.7184μs | 14.9884 KOps/s | 14.6485 KOps/s | |
| test_unbind_speed | 0.3718ms | 0.3274ms | 3.0546 KOps/s | 3.0669 KOps/s | |
| test_unbind_speed_stack0 | 0.3877ms | 0.3303ms | 3.0274 KOps/s | 3.0798 KOps/s | |
| test_unbind_speed_stack1 | 0.1042s | 0.8337ms | 1.1995 KOps/s | 1.1888 KOps/s | |
| test_split | 0.1037s | 1.2702ms | 787.2792 Ops/s | 786.1097 Ops/s | |
| test_chunk | 0.1041s | 1.2113ms | 825.5680 Ops/s | 924.9371 Ops/s | |
| test_to_cpu_blocking | 29.1362ms | 28.6235ms | 34.9363 Ops/s | 40.0692 Ops/s | |
| test_to_cpu_global_sync | 11.6288ms | 11.2858ms | 88.6070 Ops/s | 89.4632 Ops/s | |
| test_to_cpu_event_sync | 12.5525ms | 12.2533ms | 81.6107 Ops/s | 82.3870 Ops/s | |
| test_to_cpu_default | 0.1163s | 13.5734ms | 73.6734 Ops/s | 82.2807 Ops/s | |
| test_consolidate[False-None] | 4.2453ms | 4.1153ms | 242.9936 Ops/s | 219.6502 Ops/s | |
| test_consolidate[default-None] | 2.1139ms | 2.0284ms | 493.0039 Ops/s | 483.6884 Ops/s | |
| test_consolidate[reduce-overhead-None] | 2.0521ms | 1.9484ms | 513.2524 Ops/s | 511.2236 Ops/s | |
| test_consolidate_njt[False-None] | 8.6808ms | 8.4068ms | 118.9509 Ops/s | 118.6549 Ops/s | |
| test_to[False-False-None] | 2.2347ms | 2.0668ms | 483.8292 Ops/s | 476.1794 Ops/s | |
| test_to[True-False-None] | 2.1857ms | 1.9143ms | 522.3787 Ops/s | 513.1205 Ops/s | |
| test_to[within-False-None] | 6.2058ms | 6.1043ms | 163.8196 Ops/s | 163.1102 Ops/s | |
| test_to[True-default-None] | 7.8129ms | 7.5171ms | 133.0295 Ops/s | 129.6589 Ops/s | |
| test_to_njt[False-False-None] | 8.6629ms | 8.5187ms | 117.3883 Ops/s | 115.9740 Ops/s | |
| test_to_njt[True-False-None] | 7.1268ms | 6.9084ms | 144.7511 Ops/s | 140.0232 Ops/s | |
| test_to_njt[within-False-None] | 15.6610ms | 15.5015ms | 64.5098 Ops/s | 62.9410 Ops/s | |
| test_creation[device0] | 0.4028ms | 0.1167ms | 8.5662 KOps/s | 8.3004 KOps/s | |
| test_creation_from_tensor | 0.4546ms | 0.1141ms | 8.7660 KOps/s | 8.5158 KOps/s | |
| test_add_one[memmap_tensor0] | 0.3724ms | 6.5465μs | 152.7544 KOps/s | 147.2731 KOps/s | |
| test_contiguous[memmap_tensor0] | 26.8100μs | 0.6658μs | 1.5020 MOps/s | 2.1737 MOps/s | |
| test_stack[memmap_tensor0] | 20.4000μs | 4.5916μs | 217.7896 KOps/s | 219.3409 KOps/s | |
| test_memmaptd_index | 1.0186ms | 0.2662ms | 3.7572 KOps/s | 3.8583 KOps/s | |
| test_memmaptd_index_astensor | 0.5458ms | 0.3592ms | 2.7842 KOps/s | 2.8598 KOps/s | |
| test_memmaptd_index_op | 0.8432ms | 0.6004ms | 1.6656 KOps/s | 1.6600 KOps/s | |
| test_serialize_model | 0.3185s | 0.1630s | 6.1361 Ops/s | 7.3760 Ops/s | |
| test_serialize_model_pickle | 1.3498s | 1.1920s | 0.8389 Ops/s | 0.8266 Ops/s | |
| test_serialize_weights | 0.1397s | 0.1363s | 7.3389 Ops/s | 7.3675 Ops/s | |
| test_serialize_weights_returnearly | 0.4557s | 94.2220ms | 10.6132 Ops/s | 5.9223 Ops/s | |
| test_serialize_weights_pickle | 1.3647s | 1.2132s | 0.8242 Ops/s | 0.8238 Ops/s | |
| test_reshape_pytree | 0.2153ms | 33.6137μs | 29.7498 KOps/s | 29.9105 KOps/s | |
| test_reshape_td | 76.8920μs | 44.8032μs | 22.3198 KOps/s | 22.7702 KOps/s | |
| test_view_pytree | 0.2336ms | 33.3024μs | 30.0279 KOps/s | 30.9930 KOps/s | |
| test_view_td | 96.0510μs | 51.9268μs | 19.2579 KOps/s | 18.7643 KOps/s | |
| test_unbind_pytree | 0.2409ms | 36.7415μs | 27.2172 KOps/s | 27.0357 KOps/s | |
| test_unbind_td | 0.1379ms | 48.6993μs | 20.5342 KOps/s | 20.3478 KOps/s | |
| test_split_pytree | 0.2489ms | 42.9948μs | 23.2586 KOps/s | 23.7785 KOps/s | |
| test_split_td | 0.2076ms | 65.5165μs | 15.2633 KOps/s | 15.3953 KOps/s | |
| test_add_pytree | 0.1960ms | 42.3189μs | 23.6301 KOps/s | 23.7199 KOps/s | |
| test_add_td | 93.8920μs | 53.1536μs | 18.8134 KOps/s | 18.2957 KOps/s | |
| test_compile_add_one_nested[tensordict-compile] | 0.1940ms | 0.1416ms | 7.0626 KOps/s | 6.4524 KOps/s | |
| test_compile_add_one_nested[tensordict-eager] | 0.3121ms | 0.1951ms | 5.1269 KOps/s | 5.1849 KOps/s | |
| test_compile_add_one_nested[pytree-compile] | 0.1584ms | 0.1095ms | 9.1306 KOps/s | 8.9889 KOps/s | |
| test_compile_add_one_nested[pytree-eager] | 0.4469ms | 0.1821ms | 5.4912 KOps/s | 5.3821 KOps/s | |
| test_compile_copy_nested[tensordict-compile] | 76.2410μs | 29.6336μs | 33.7454 KOps/s | 30.4929 KOps/s | |
| test_compile_copy_nested[tensordict-eager] | 84.9820μs | 52.9126μs | 18.8991 KOps/s | 18.9621 KOps/s | |
| test_compile_copy_nested[pytree-compile] | 39.8410μs | 9.9477μs | 100.5257 KOps/s | 101.8539 KOps/s | |
| test_compile_copy_nested[pytree-eager] | 0.4615ms | 70.4838μs | 14.1877 KOps/s | 14.6012 KOps/s | |
| test_compile_add_one_flat[tensordict-compile] | 0.2301ms | 0.1786ms | 5.5994 KOps/s | 5.3595 KOps/s | |
| test_compile_add_one_flat[tensordict-eager] | 6.2944ms | 0.2579ms | 3.8767 KOps/s | 3.8938 KOps/s | |
| test_compile_add_one_flat[tensorclass-compile] | 0.1557ms | 0.1179ms | 8.4786 KOps/s | 8.3221 KOps/s | |
| test_compile_add_one_flat[tensorclass-eager] | 0.1055ms | 68.9842μs | 14.4961 KOps/s | 14.3627 KOps/s | |
| test_compile_add_one_flat[pytree-compile] | 0.1966ms | 0.1595ms | 6.2701 KOps/s | 6.1121 KOps/s | |
| test_compile_add_one_flat[pytree-eager] | 0.8337ms | 0.5342ms | 1.8720 KOps/s | 1.7855 KOps/s | |
| test_compile_add_self_flat[tensordict-eager] | 0.3638ms | 0.3104ms | 3.2218 KOps/s | 3.2077 KOps/s | |
| test_compile_add_self_flat[tensordict-compile] | 0.2342ms | 0.1804ms | 5.5441 KOps/s | 4.9742 KOps/s | |
| test_compile_add_self_flat[tensorclass-eager] | 0.1250ms | 84.7258μs | 11.8028 KOps/s | 11.7279 KOps/s | |
| test_compile_add_self_flat[tensorclass-compile] | 0.1590ms | 0.1211ms | 8.2591 KOps/s | 7.9143 KOps/s | |
| test_compile_add_self_flat[pytree-eager] | 0.6800ms | 0.4417ms | 2.2642 KOps/s | 2.1575 KOps/s | |
| test_compile_add_self_flat[pytree-compile] | 0.2026ms | 0.1594ms | 6.2729 KOps/s | 5.7894 KOps/s | |
| test_compile_copy_flat[tensordict-compile] | 60.1010μs | 24.2722μs | 41.1993 KOps/s | 38.3630 KOps/s | |
| test_compile_copy_flat[tensordict-eager] | 69.1010μs | 41.9963μs | 23.8116 KOps/s | 24.0958 KOps/s | |
| test_compile_copy_flat[pytree-compile] | 46.7610μs | 10.8535μs | 92.1362 KOps/s | 92.4748 KOps/s | |
| test_compile_copy_flat[pytree-eager] | 0.1807s | 62.0356μs | 16.1198 KOps/s | 19.2389 KOps/s | |
| test_compile_assign_and_add[tensordict-compile] | 2.0426ms | 0.1798ms | 5.5621 KOps/s | 5.2644 KOps/s | |
| test_compile_assign_and_add[tensordict-eager] | 3.4550ms | 3.2978ms | 303.2331 Ops/s | 301.3114 Ops/s | |
| test_compile_assign_and_add[pytree-compile] | 1.9902ms | 0.1625ms | 6.1531 KOps/s | 6.1412 KOps/s | |
| test_compile_assign_and_add[pytree-eager] | 2.9409ms | 2.8168ms | 355.0068 Ops/s | 350.6682 Ops/s | |
| test_compile_indexing[tensor-tensordict-compile] | 0.1712ms | 0.1096ms | 9.1221 KOps/s | 8.9213 KOps/s | |
| test_compile_indexing[tensor-tensordict-eager] | 0.3651ms | 72.5013μs | 13.7929 KOps/s | 13.8126 KOps/s | |
| test_compile_indexing[tensor-tensorclass-compile] | 0.1368ms | 97.4024μs | 10.2667 KOps/s | 10.2187 KOps/s | |
| test_compile_indexing[tensor-tensorclass-eager] | 0.2708ms | 44.4807μs | 22.4817 KOps/s | 22.2434 KOps/s | |
| test_compile_indexing[tensor-pytree-compile] | 0.1452ms | 98.0281μs | 10.2012 KOps/s | 10.1776 KOps/s | |
| test_compile_indexing[tensor-pytree-eager] | 0.2597ms | 44.6191μs | 22.4119 KOps/s | 22.2359 KOps/s | |
| test_compile_indexing[slice-tensordict-compile] | 0.1279ms | 56.0741μs | 17.8336 KOps/s | 16.9567 KOps/s | |
| test_compile_indexing[slice-tensordict-eager] | 0.2254ms | 27.6765μs | 36.1317 KOps/s | 36.3065 KOps/s | |
| test_compile_indexing[slice-tensorclass-compile] | 82.4520μs | 45.9461μs | 21.7646 KOps/s | 22.1038 KOps/s | |
| test_compile_indexing[slice-tensorclass-eager] | 0.2601ms | 22.8239μs | 43.8137 KOps/s | 43.9644 KOps/s | |
| test_compile_indexing[slice-pytree-compile] | 86.9220μs | 46.0157μs | 21.7317 KOps/s | 21.8590 KOps/s | |
| test_compile_indexing[slice-pytree-eager] | 0.2696ms | 22.6049μs | 44.2381 KOps/s | 44.2704 KOps/s | |
| test_compile_indexing[int-tensordict-compile] | 0.1073ms | 56.9620μs | 17.5556 KOps/s | 16.8742 KOps/s | |
| test_compile_indexing[int-tensordict-eager] | 0.2304ms | 27.5857μs | 36.2507 KOps/s | 36.1035 KOps/s | |
| test_compile_indexing[int-tensorclass-compile] | 93.0420μs | 45.4771μs | 21.9891 KOps/s | 21.7999 KOps/s | |
| test_compile_indexing[int-tensorclass-eager] | 0.2617ms | 22.7833μs | 43.8918 KOps/s | 44.1711 KOps/s | |
| test_compile_indexing[int-pytree-compile] | 92.7320μs | 45.6473μs | 21.9071 KOps/s | 21.7277 KOps/s | |
| test_compile_indexing[int-pytree-eager] | 0.2611ms | 22.1474μs | 45.1520 KOps/s | 44.1670 KOps/s | |
| test_mod_add[eager] | 91.3210μs | 48.8463μs | 20.4724 KOps/s | 19.9961 KOps/s | |
| test_mod_add[compile] | 0.2098ms | 0.1055ms | 9.4815 KOps/s | 9.4581 KOps/s | |
| test_mod_add[compile-overhead] | 0.2355ms | 0.1499ms | 6.6695 KOps/s | 6.5438 KOps/s | |
| test_mod_wrap[eager] | 0.3616ms | 0.2901ms | 3.4475 KOps/s | 3.3618 KOps/s | |
| test_mod_wrap[compile] | 0.4528ms | 0.3612ms | 2.7689 KOps/s | 2.8293 KOps/s | |
| test_mod_wrap[compile-overhead] | 7.1843ms | 3.9842ms | 250.9898 Ops/s | 247.7151 Ops/s | |
| test_mod_wrap_and_backward[eager] | 1.6524ms | 1.5273ms | 654.7638 Ops/s | 646.8489 Ops/s | |
| test_mod_wrap_and_backward[compile] | 1.6436ms | 1.4504ms | 689.4617 Ops/s | 686.4058 Ops/s | |
| test_mod_wrap_and_backward[compile-overhead] | 1.2579ms | 0.9570ms | 1.0449 KOps/s | 1.1021 KOps/s | |
| test_seq_add[eager] | 0.2586ms | 0.1539ms | 6.4984 KOps/s | 6.3409 KOps/s | |
| test_seq_add[compile] | 0.4629ms | 0.1179ms | 8.4791 KOps/s | 8.3965 KOps/s | |
| test_seq_add[compile-overhead] | 0.1987ms | 0.1564ms | 6.3945 KOps/s | 6.2763 KOps/s | |
| test_seq_wrap[eager] | 0.7947ms | 0.5475ms | 1.8265 KOps/s | 1.8184 KOps/s | |
| test_seq_wrap[compile] | 0.4159ms | 0.3664ms | 2.7295 KOps/s | 2.6170 KOps/s | |
| test_seq_wrap[compile-overhead] | 0.3425ms | 0.2666ms | 3.7511 KOps/s | 3.7069 KOps/s | |
| test_func_call_runtime[False-eager] | 0.9764ms | 0.8846ms | 1.1305 KOps/s | 1.1597 KOps/s | |
| test_func_call_runtime[False-compile] | 1.0121ms | 0.9092ms | 1.0998 KOps/s | 1.0833 KOps/s | |
| test_func_call_runtime[False-compile-overhead] | 0.5319ms | 0.4602ms | 2.1728 KOps/s | 2.1262 KOps/s | |
| test_func_call_runtime[True-eager] | 1.1925ms | 1.0922ms | 915.6017 Ops/s | 908.6231 Ops/s | |
| test_func_call_runtime[True-compile] | 0.9852ms | 0.9197ms | 1.0873 KOps/s | 1.0646 KOps/s | |
| test_func_call_runtime[True-compile-overhead] | 0.5237ms | 0.4737ms | 2.1112 KOps/s | 2.0567 KOps/s | |
| test_func_call_cm_runtime[False-eager] | 0.9136ms | 0.8388ms | 1.1922 KOps/s | 1.1121 KOps/s | |
| test_func_call_cm_runtime[False-compile] | 1.1313ms | 0.9136ms | 1.0945 KOps/s | 1.0707 KOps/s | |
| test_func_call_cm_runtime[False-compile-overhead] | 0.5232ms | 0.4635ms | 2.1575 KOps/s | 2.1150 KOps/s | |
| test_func_call_cm_runtime[True-eager] | 1.4306ms | 1.2332ms | 810.9048 Ops/s | 803.0255 Ops/s | |
| test_func_call_cm_runtime[True-compile] | 1.2216ms | 0.9662ms | 1.0350 KOps/s | 1.0307 KOps/s | |
| test_func_call_cm_runtime[True-compile-overhead] | 0.5855ms | 0.5088ms | 1.9655 KOps/s | 1.9280 KOps/s | |
| test_vmap_func_call_cm_runtime[eager] | 2.8750ms | 2.3656ms | 422.7308 Ops/s | 413.9631 Ops/s | |
| test_vmap_func_call_cm_runtime[compile] | 1.1242ms | 0.9833ms | 1.0170 KOps/s | 1.0007 KOps/s | |
| test_vmap_func_call_cm_runtime[compile-overhead] | 0.6301ms | 0.5132ms | 1.9486 KOps/s | 1.9061 KOps/s | |
| test_distributed | 3.0330ms | 0.1642ms | 6.0890 KOps/s | 6.0613 KOps/s | |
| test_tdmodule | 68.8010μs | 27.9745μs | 35.7468 KOps/s | 34.2219 KOps/s | |
| test_tdmodule_dispatch | 76.4320μs | 45.1153μs | 22.1654 KOps/s | 21.9300 KOps/s | |
| test_tdseq | 47.7010μs | 27.6821μs | 36.1245 KOps/s | 35.8260 KOps/s | |
| test_tdseq_dispatch | 65.7120μs | 47.6163μs | 21.0012 KOps/s | 20.8077 KOps/s | |
| test_instantiation_functorch | 2.1974ms | 2.0802ms | 480.7224 Ops/s | 480.0977 Ops/s | |
| test_exec_functorch | 0.2506ms | 0.1816ms | 5.5078 KOps/s | 5.4880 KOps/s | |
| test_exec_functional_call | 0.2909ms | 0.1682ms | 5.9445 KOps/s | 6.1415 KOps/s | |
| test_exec_td_decorator | 0.4441ms | 0.2400ms | 4.1667 KOps/s | 4.2294 KOps/s | |
| test_vmap_mlp_speed_decorator[True-True] | 0.9957ms | 0.8245ms | 1.2129 KOps/s | 1.2015 KOps/s | |
| test_vmap_mlp_speed_decorator[True-False] | 1.0538ms | 0.8235ms | 1.2144 KOps/s | 1.1984 KOps/s | |
| test_vmap_mlp_speed_decorator[False-True] | 0.9233ms | 0.7172ms | 1.3943 KOps/s | 1.3696 KOps/s | |
| test_vmap_mlp_speed_decorator[False-False] | 0.9032ms | 0.7131ms | 1.4024 KOps/s | 1.3705 KOps/s | |
| test_vmap_transformer_speed_decorator[True-True] | 20.8543ms | 20.6672ms | 48.3858 Ops/s | 47.8091 Ops/s | |
| test_vmap_transformer_speed_decorator[True-False] | 20.9638ms | 20.7190ms | 48.2649 Ops/s | 47.6908 Ops/s | |
| test_vmap_transformer_speed_decorator[False-True] | 21.1460ms | 20.4988ms | 48.7833 Ops/s | 48.3007 Ops/s | |
| test_vmap_transformer_speed_decorator[False-False] | 21.3641ms | 20.6266ms | 48.4811 Ops/s | 48.4059 Ops/s | |
| test_to_module_speed[True] | 2.0357ms | 1.4709ms | 679.8604 Ops/s | 676.8233 Ops/s | |
| test_to_module_speed[False] | 2.0130ms | 1.4513ms | 689.0379 Ops/s | 685.2291 Ops/s | |
| test_tc_init | 68.3910μs | 45.5391μs | 21.9592 KOps/s | 21.5674 KOps/s | |
| test_tc_init_tensor_only | 35.6210μs | 9.7951μs | 102.0919 KOps/s | 99.8962 KOps/s | |
| test_tc_init_nested | 0.1407ms | 90.0015μs | 11.1109 KOps/s | 11.0029 KOps/s | |
| test_tc_init_many_fields | 84.5320μs | 16.4160μs | 60.9163 KOps/s | 58.5543 KOps/s | |
| test_tc_first_layer_tensor | 18.3700μs | 1.8571μs | 538.4666 KOps/s | 532.0797 KOps/s | |
| test_tc_first_layer_tensor_only | 5.4701μs | 0.7447μs | 1.3427 MOps/s | 1.3328 MOps/s | |
| test_tc_first_layer_tensor_set | 37.9010μs | 3.9787μs | 251.3369 KOps/s | 252.6251 KOps/s | |
| test_tc_first_layer_tensor_only_set | 24.3710μs | 3.2087μs | 311.6560 KOps/s | 309.7789 KOps/s | |
| test_tc_first_layer_nontensor | 37.4200μs | 6.2621μs | 159.6907 KOps/s | 157.8029 KOps/s | |
| test_tc_second_layer_tensor | 34.6600μs | 4.4392μs | 225.2683 KOps/s | 226.3248 KOps/s | |
| test_tc_second_layer_nontensor | 41.7610μs | 8.8626μs | 112.8342 KOps/s | 110.9538 KOps/s | |
| test_unbind | 0.2602s | 16.3130ms | 61.3008 Ops/s | 63.2965 Ops/s | |
| test_full_like | 5.5392ms | 4.3759ms | 228.5231 Ops/s | 228.7359 Ops/s | |
| test_zeros_like | 4.4978ms | 4.3630ms | 229.1978 Ops/s | 228.4111 Ops/s | |
| test_ones_like | 4.5006ms | 4.3699ms | 228.8381 Ops/s | 228.8494 Ops/s | |
| test_clone | 6.5104ms | 6.4108ms | 155.9859 Ops/s | 154.7485 Ops/s | |
| test_squeeze | 0.2002ms | 14.5616μs | 68.6737 KOps/s | 68.6058 KOps/s | |
| test_unsqueeze | 0.2688ms | 0.1101ms | 9.0827 KOps/s | 8.9751 KOps/s | |
| test_split | 0.2373ms | 0.1858ms | 5.3819 KOps/s | 5.3879 KOps/s | |
| test_permute | 0.3183ms | 0.2111ms | 4.7380 KOps/s | 4.8883 KOps/s | |
| test_stack | 51.6150ms | 50.8354ms | 19.6713 Ops/s | 19.5856 Ops/s | |
| test_cat | 51.5173ms | 51.2171ms | 19.5247 Ops/s | 19.5567 Ops/s | |
| test_sequential_tensordict | 0.3314ms | 0.2206ms | 4.5334 KOps/s | 4.4131 KOps/s | |
| test_sequential_graph_module | 0.5100ms | 0.1170ms | 8.5478 KOps/s | 7.9694 KOps/s | |
| test_nested_tensordict | 0.4158ms | 0.2796ms | 3.5761 KOps/s | 3.4851 KOps/s | |
| test_nested_graph_module | 0.5400ms | 0.1264ms | 7.9135 KOps/s | 7.3327 KOps/s |
|
| 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 | |
| test_plain_set_stack_nested | 45.5910μs | 14.9035μs | 67.0981 KOps/s | 67.2045 KOps/s | |
| test_plain_set_nested_inplace | 36.6210μs | 16.4157μs | 60.9174 KOps/s | 59.2973 KOps/s | |
| test_plain_set_stack_nested_inplace | 42.1600μs | 16.3016μs | 61.3435 KOps/s | 60.4733 KOps/s | |
| test_items | 32.2610μs | 5.8482μs | 170.9934 KOps/s | 168.8200 KOps/s | |
| test_items_nested | 0.4865ms | 0.4547ms | 2.1990 KOps/s | 2.1983 KOps/s | |
| test_items_nested_locked | 0.5223ms | 0.4574ms | 2.1864 KOps/s | 2.1855 KOps/s | |
| test_items_nested_leaf | 0.1345ms | 96.6506μs | 10.3465 KOps/s | 10.4761 KOps/s | |
| test_items_stack_nested | 0.4892ms | 0.4546ms | 2.1995 KOps/s | 2.1943 KOps/s | |
| test_items_stack_nested_leaf | 0.1235ms | 95.8758μs | 10.4302 KOps/s | 10.2397 KOps/s | |
| test_items_stack_nested_locked | 0.5031ms | 0.4568ms | 2.1891 KOps/s | 2.1795 KOps/s | |
| test_keys | 32.2410μs | 4.3939μs | 227.5858 KOps/s | 238.5046 KOps/s | |
| test_keys_nested | 0.1500ms | 0.1194ms | 8.3768 KOps/s | 8.2991 KOps/s | |
| test_keys_nested_locked | 2.2264ms | 0.1289ms | 7.7574 KOps/s | 7.0254 KOps/s | |
| test_keys_nested_leaf | 0.4273ms | 0.1095ms | 9.1364 KOps/s | 8.8941 KOps/s | |
| test_keys_stack_nested | 0.2768ms | 0.1191ms | 8.3970 KOps/s | 8.1705 KOps/s | |
| test_keys_stack_nested_leaf | 0.1501ms | 0.1094ms | 9.1376 KOps/s | 8.9332 KOps/s | |
| test_keys_stack_nested_locked | 0.1579ms | 0.1279ms | 7.8188 KOps/s | 7.6637 KOps/s | |
| test_values | 5.8800μs | 1.0265μs | 974.1928 KOps/s | 970.7254 KOps/s | |
| test_values_nested | 72.0510μs | 47.5578μs | 21.0271 KOps/s | 20.7164 KOps/s | |
| test_values_nested_locked | 74.7610μs | 50.7543μs | 19.7028 KOps/s | 19.5340 KOps/s | |
| test_values_nested_leaf | 76.6910μs | 53.6596μs | 18.6360 KOps/s | 18.2046 KOps/s | |
| test_values_stack_nested | 74.2000μs | 47.4325μs | 21.0826 KOps/s | 20.5257 KOps/s | |
| test_values_stack_nested_leaf | 82.8510μs | 54.2405μs | 18.4364 KOps/s | 18.1921 KOps/s | |
| test_values_stack_nested_locked | 80.8610μs | 50.5523μs | 19.7815 KOps/s | 19.4720 KOps/s | |
| test_membership | 5.7133μs | 0.8752μs | 1.1426 MOps/s | 1.1638 MOps/s | |
| test_membership_nested | 25.6410μs | 2.9243μs | 341.9643 KOps/s | 346.7291 KOps/s | |
| test_membership_nested_leaf | 34.6600μs | 2.9232μs | 342.0959 KOps/s | 348.4674 KOps/s | |
| test_membership_stacked_nested | 23.7110μs | 2.9144μs | 343.1208 KOps/s | 342.4021 KOps/s | |
| test_membership_stacked_nested_leaf | 23.8600μs | 2.9275μs | 341.5926 KOps/s | 344.1591 KOps/s | |
| test_membership_nested_last | 35.9100μs | 4.3313μs | 230.8770 KOps/s | 229.8569 KOps/s | |
| test_membership_nested_leaf_last | 46.7410μs | 4.2972μs | 232.7094 KOps/s | 228.8224 KOps/s | |
| test_membership_stacked_nested_last | 22.6300μs | 4.3340μs | 230.7357 KOps/s | 230.7970 KOps/s | |
| test_membership_stacked_nested_leaf_last | 25.0800μs | 4.3458μs | 230.1093 KOps/s | 234.0429 KOps/s | |
| test_nested_getleaf | 80.1210μs | 21.3191μs | 46.9064 KOps/s | 47.5993 KOps/s | |
| test_nested_get | 44.6300μs | 20.2624μs | 49.3525 KOps/s | 49.4813 KOps/s | |
| test_stacked_getleaf | 50.9310μs | 21.6704μs | 46.1458 KOps/s | 46.5277 KOps/s | |
| test_stacked_get | 50.5510μs | 20.5217μs | 48.7290 KOps/s | 49.6260 KOps/s | |
| test_nested_getitemleaf | 42.1400μs | 21.8484μs | 45.7700 KOps/s | 46.2110 KOps/s | |
| test_nested_getitem | 48.7010μs | 20.8533μs | 47.9540 KOps/s | 48.1775 KOps/s | |
| test_stacked_getitemleaf | 58.1710μs | 21.9363μs | 45.5865 KOps/s | 45.9048 KOps/s | |
| test_stacked_getitem | 44.5700μs | 20.9402μs | 47.7551 KOps/s | 47.9593 KOps/s | |
| test_lock_nested | 0.6420ms | 0.4673ms | 2.1400 KOps/s | 2.0868 KOps/s | |
| test_lock_stack_nested | 0.5219ms | 0.4747ms | 2.1067 KOps/s | 2.0602 KOps/s | |
| test_unlock_nested | 0.5365ms | 0.3800ms | 2.6316 KOps/s | 2.6154 KOps/s | |
| test_unlock_stack_nested | 0.4055ms | 0.3812ms | 2.6236 KOps/s | 2.5645 KOps/s | |
| test_flatten_speed | 0.1578ms | 0.1228ms | 8.1411 KOps/s | 8.0898 KOps/s | |
| test_unflatten_speed | 0.6045ms | 0.5713ms | 1.7503 KOps/s | 1.7377 KOps/s | |
| test_common_ops | 0.8189ms | 0.6758ms | 1.4797 KOps/s | 1.4506 KOps/s | |
| test_creation | 0.1224ms | 2.9236μs | 342.0435 KOps/s | 339.9941 KOps/s | |
| test_creation_empty | 31.1600μs | 6.1693μs | 162.0937 KOps/s | 161.3743 KOps/s | |
| test_creation_nested_1 | 32.8500μs | 10.4826μs | 95.3958 KOps/s | 95.2035 KOps/s | |
| test_creation_nested_2 | 37.8800μs | 11.9302μs | 83.8207 KOps/s | 84.0326 KOps/s | |
| test_creation_many_keys[10] | 46.1300μs | 18.3189μs | 54.5885 KOps/s | 54.7449 KOps/s | |
| test_creation_many_keys[50] | 0.1012ms | 78.3686μs | 12.7602 KOps/s | 12.7518 KOps/s | |
| test_creation_many_keys[100] | 0.1924ms | 0.1541ms | 6.4875 KOps/s | 6.5927 KOps/s | |
| test_creation_nested_many_keys[10] | 75.0110μs | 39.6116μs | 25.2451 KOps/s | 25.4172 KOps/s | |
| test_creation_nested_many_keys[50] | 0.1987ms | 0.1608ms | 6.2182 KOps/s | 6.2920 KOps/s | |
| test_clone | 53.2010μs | 13.1141μs | 76.2539 KOps/s | 73.3316 KOps/s | |
| test_getitem[int] | 1.6699ms | 14.5137μs | 68.9003 KOps/s | 55.2406 KOps/s | |
| test_getitem[slice_int] | 0.1393ms | 24.5726μs | 40.6957 KOps/s | 40.2487 KOps/s | |
| test_getitem[range] | 0.1759ms | 62.7166μs | 15.9447 KOps/s | 15.8656 KOps/s | |
| test_getitem[tuple] | 0.1426ms | 23.8964μs | 41.8473 KOps/s | 40.9576 KOps/s | |
| test_getitem[list] | 0.1872ms | 57.6451μs | 17.3475 KOps/s | 16.9944 KOps/s | |
| test_setitem_dim[int] | 59.2000μs | 26.0464μs | 38.3931 KOps/s | 37.6567 KOps/s | |
| test_setitem_dim[slice_int] | 65.0010μs | 43.7053μs | 22.8805 KOps/s | 22.8543 KOps/s | |
| test_setitem_dim[range] | 0.1173ms | 94.1894μs | 10.6169 KOps/s | 10.5063 KOps/s | |
| test_setitem_dim[tuple] | 65.8010μs | 40.2951μs | 24.8169 KOps/s | 24.4182 KOps/s | |
| test_setitem | 71.9700μs | 17.6751μs | 56.5769 KOps/s | 54.5869 KOps/s | |
| test_set | 61.8410μs | 16.6793μs | 59.9546 KOps/s | 57.5105 KOps/s | |
| test_set_shared | 0.5664ms | 0.2040ms | 4.9012 KOps/s | 4.8597 KOps/s | |
| test_update | 0.2171ms | 21.0206μs | 47.5723 KOps/s | 45.5298 KOps/s | |
| test_update_nested | 75.3310μs | 33.4101μs | 29.9311 KOps/s | 29.4403 KOps/s | |
| test_update__nested | 0.4594ms | 33.6717μs | 29.6985 KOps/s | 27.8916 KOps/s | |
| test_set_nested | 60.2100μs | 18.4219μs | 54.2831 KOps/s | 52.4798 KOps/s | |
| test_set_nested_new | 65.7600μs | 23.4964μs | 42.5598 KOps/s | 41.6871 KOps/s | |
| test_select | 71.8310μs | 40.4687μs | 24.7104 KOps/s | 24.1407 KOps/s | |
| test_select_nested | 0.1086ms | 74.9644μs | 13.3397 KOps/s | 13.3804 KOps/s | |
| test_exclude_nested | 0.1205ms | 92.6055μs | 10.7985 KOps/s | 10.8050 KOps/s | |
| test_empty[True] | 0.4431ms | 0.3890ms | 2.5704 KOps/s | 2.5464 KOps/s | |
| test_empty[False] | 7.6952μs | 1.3128μs | 761.7101 KOps/s | 752.0084 KOps/s | |
| test_to | 0.1027ms | 71.4021μs | 14.0052 KOps/s | 13.6464 KOps/s | |
| test_to_nonblocking | 0.1008ms | 64.7811μs | 15.4366 KOps/s | 15.4285 KOps/s | |
| test_unbind_speed | 0.3639ms | 0.3253ms | 3.0743 KOps/s | 3.0294 KOps/s | |
| test_unbind_speed_stack0 | 0.3715ms | 0.3244ms | 3.0824 KOps/s | 3.0717 KOps/s | |
| test_unbind_speed_stack1 | 0.1045s | 0.8292ms | 1.2060 KOps/s | 1.1705 KOps/s | |
| test_split | 0.1045s | 1.2716ms | 786.4397 Ops/s | 777.3484 Ops/s | |
| test_chunk | 0.1040s | 1.2077ms | 828.0387 Ops/s | 915.1694 Ops/s | |
| test_to_cpu_blocking | 19.1893ms | 19.0764ms | 52.4207 Ops/s | 31.5390 Ops/s | |
| test_to_cpu_global_sync | 11.1104ms | 10.9896ms | 90.9955 Ops/s | 88.8616 Ops/s | |
| test_to_cpu_event_sync | 12.2831ms | 12.0066ms | 83.2875 Ops/s | 81.6599 Ops/s | |
| test_to_cpu_default | 0.1162s | 13.2404ms | 75.5262 Ops/s | 81.6999 Ops/s | |
| test_consolidate[False-None] | 4.3047ms | 4.1481ms | 241.0762 Ops/s | 216.7138 Ops/s | |
| test_consolidate[default-None] | 2.1748ms | 2.0470ms | 488.5201 Ops/s | 483.4830 Ops/s | |
| test_consolidate[reduce-overhead-None] | 2.1005ms | 1.9634ms | 509.3105 Ops/s | 500.7562 Ops/s | |
| test_consolidate_njt[False-None] | 8.6566ms | 8.4623ms | 118.1714 Ops/s | 116.5662 Ops/s | |
| test_to[False-False-None] | 2.1532ms | 2.0467ms | 488.5971 Ops/s | 480.9401 Ops/s | |
| test_to[True-False-None] | 2.1986ms | 1.9457ms | 513.9577 Ops/s | 525.8979 Ops/s | |
| test_to[within-False-None] | 6.2895ms | 6.1573ms | 162.4094 Ops/s | 161.7854 Ops/s | |
| test_to[True-default-None] | 7.6726ms | 7.5595ms | 132.2847 Ops/s | 130.0377 Ops/s | |
| test_to_njt[False-False-None] | 8.7126ms | 8.5342ms | 117.1761 Ops/s | 115.5047 Ops/s | |
| test_to_njt[True-False-None] | 7.1689ms | 6.9997ms | 142.8624 Ops/s | 141.4719 Ops/s | |
| test_to_njt[within-False-None] | 15.9317ms | 15.7642ms | 63.4347 Ops/s | 63.1094 Ops/s | |
| test_creation[device0] | 0.4469ms | 0.1174ms | 8.5172 KOps/s | 8.1877 KOps/s | |
| test_creation_from_tensor | 0.4568ms | 0.1139ms | 8.7775 KOps/s | 8.4509 KOps/s | |
| test_add_one[memmap_tensor0] | 0.4256ms | 6.4851μs | 154.1997 KOps/s | 149.3105 KOps/s | |
| test_contiguous[memmap_tensor0] | 31.4410μs | 0.6686μs | 1.4956 MOps/s | 2.2275 MOps/s | |
| test_stack[memmap_tensor0] | 57.1200μs | 4.6354μs | 215.7306 KOps/s | 209.0810 KOps/s | |
| test_memmaptd_index | 1.0860ms | 0.2663ms | 3.7559 KOps/s | 3.7287 KOps/s | |
| test_memmaptd_index_astensor | 0.5118ms | 0.3581ms | 2.7926 KOps/s | 2.7464 KOps/s | |
| test_memmaptd_index_op | 0.8725ms | 0.6044ms | 1.6545 KOps/s | 1.6309 KOps/s | |
| test_serialize_model | 0.1387s | 0.1364s | 7.3322 Ops/s | 7.3484 Ops/s | |
| test_serialize_model_pickle | 1.3472s | 1.1926s | 0.8385 Ops/s | 0.8260 Ops/s | |
| test_serialize_weights | 0.1408s | 0.1347s | 7.4266 Ops/s | 7.3939 Ops/s | |
| test_serialize_weights_returnearly | 0.3757s | 93.9065ms | 10.6489 Ops/s | 14.6216 Ops/s | |
| test_serialize_weights_pickle | 1.3657s | 1.1883s | 0.8415 Ops/s | 0.8245 Ops/s | |
| test_reshape_pytree | 0.2133ms | 33.7608μs | 29.6202 KOps/s | 29.5971 KOps/s | |
| test_reshape_td | 73.1110μs | 44.2177μs | 22.6154 KOps/s | 22.2993 KOps/s | |
| test_view_pytree | 0.2526ms | 33.4013μs | 29.9389 KOps/s | 29.6836 KOps/s | |
| test_view_td | 82.1710μs | 51.8336μs | 19.2925 KOps/s | 18.7960 KOps/s | |
| test_unbind_pytree | 0.2366ms | 37.1554μs | 26.9140 KOps/s | 26.4851 KOps/s | |
| test_unbind_td | 0.1174ms | 49.6290μs | 20.1495 KOps/s | 20.3667 KOps/s | |
| test_split_pytree | 0.2556ms | 42.8721μs | 23.3252 KOps/s | 23.0854 KOps/s | |
| test_split_td | 0.1343ms | 64.3527μs | 15.5394 KOps/s | 15.4168 KOps/s | |
| test_add_pytree | 0.2311ms | 42.5488μs | 23.5024 KOps/s | 22.8257 KOps/s | |
| test_add_td | 92.3620μs | 53.0258μs | 18.8587 KOps/s | 17.9073 KOps/s | |
| test_compile_add_one_nested[tensordict-compile] | 0.2079ms | 0.1385ms | 7.2198 KOps/s | 6.9448 KOps/s | |
| test_compile_add_one_nested[tensordict-eager] | 0.3140ms | 0.1940ms | 5.1551 KOps/s | 5.2684 KOps/s | |
| test_compile_add_one_nested[pytree-compile] | 0.1763ms | 0.1081ms | 9.2511 KOps/s | 8.9448 KOps/s | |
| test_compile_add_one_nested[pytree-eager] | 0.4766ms | 0.1842ms | 5.4295 KOps/s | 5.4278 KOps/s | |
| test_compile_copy_nested[tensordict-compile] | 70.6510μs | 30.6021μs | 32.6775 KOps/s | 33.0887 KOps/s | |
| test_compile_copy_nested[tensordict-eager] | 93.6610μs | 52.6545μs | 18.9917 KOps/s | 19.0549 KOps/s | |
| test_compile_copy_nested[pytree-compile] | 44.4410μs | 9.9010μs | 101.0000 KOps/s | 100.2981 KOps/s | |
| test_compile_copy_nested[pytree-eager] | 0.4279ms | 70.4041μs | 14.2037 KOps/s | 14.0334 KOps/s | |
| test_compile_add_one_flat[tensordict-compile] | 0.2320ms | 0.1772ms | 5.6434 KOps/s | 5.1738 KOps/s | |
| test_compile_add_one_flat[tensordict-eager] | 0.3756ms | 0.2529ms | 3.9542 KOps/s | 3.8745 KOps/s | |
| test_compile_add_one_flat[tensorclass-compile] | 0.1653ms | 0.1159ms | 8.6290 KOps/s | 8.0431 KOps/s | |
| test_compile_add_one_flat[tensorclass-eager] | 0.1258ms | 68.8639μs | 14.5214 KOps/s | 14.3198 KOps/s | |
| test_compile_add_one_flat[pytree-compile] | 0.2165ms | 0.1588ms | 6.2985 KOps/s | 6.0294 KOps/s | |
| test_compile_add_one_flat[pytree-eager] | 0.8526ms | 0.5262ms | 1.9005 KOps/s | 1.8270 KOps/s | |
| test_compile_add_self_flat[tensordict-eager] | 0.3547ms | 0.3058ms | 3.2703 KOps/s | 3.1821 KOps/s | |
| test_compile_add_self_flat[tensordict-compile] | 0.3425ms | 0.1859ms | 5.3786 KOps/s | 5.0678 KOps/s | |
| test_compile_add_self_flat[tensorclass-eager] | 0.2794ms | 85.4661μs | 11.7005 KOps/s | 11.8334 KOps/s | |
| test_compile_add_self_flat[tensorclass-compile] | 0.1600ms | 0.1175ms | 8.5087 KOps/s | 7.7757 KOps/s | |
| test_compile_add_self_flat[pytree-eager] | 0.6598ms | 0.4366ms | 2.2902 KOps/s | 2.2264 KOps/s | |
| test_compile_add_self_flat[pytree-compile] | 0.2154ms | 0.1583ms | 6.3168 KOps/s | 6.0378 KOps/s | |
| test_compile_copy_flat[tensordict-compile] | 91.4010μs | 24.4982μs | 40.8193 KOps/s | 37.0589 KOps/s | |
| test_compile_copy_flat[tensordict-eager] | 82.7710μs | 41.6308μs | 24.0207 KOps/s | 23.9978 KOps/s | |
| test_compile_copy_flat[pytree-compile] | 45.1910μs | 10.8997μs | 91.7457 KOps/s | 91.2452 KOps/s | |
| test_compile_copy_flat[pytree-eager] | 0.3862ms | 52.2609μs | 19.1348 KOps/s | 19.2242 KOps/s | |
| test_compile_assign_and_add[tensordict-compile] | 2.0334ms | 0.1736ms | 5.7598 KOps/s | 5.3468 KOps/s | |
| test_compile_assign_and_add[tensordict-eager] | 3.5163ms | 3.2988ms | 303.1371 Ops/s | 297.8432 Ops/s | |
| test_compile_assign_and_add[pytree-compile] | 2.0038ms | 0.1619ms | 6.1754 KOps/s | 5.8403 KOps/s | |
| test_compile_assign_and_add[pytree-eager] | 2.9075ms | 2.7651ms | 361.6549 Ops/s | 353.4527 Ops/s | |
| test_compile_indexing[tensor-tensordict-compile] | 0.1685ms | 0.1117ms | 8.9561 KOps/s | 8.7543 KOps/s | |
| test_compile_indexing[tensor-tensordict-eager] | 0.3238ms | 76.5610μs | 13.0615 KOps/s | 13.3347 KOps/s | |
| test_compile_indexing[tensor-tensorclass-compile] | 0.1431ms | 96.3250μs | 10.3815 KOps/s | 10.0575 KOps/s | |
| test_compile_indexing[tensor-tensorclass-eager] | 0.2640ms | 47.1551μs | 21.2066 KOps/s | 21.4620 KOps/s | |
| test_compile_indexing[tensor-pytree-compile] | 0.1449ms | 96.6449μs | 10.3472 KOps/s | 10.0948 KOps/s | |
| test_compile_indexing[tensor-pytree-eager] | 0.2629ms | 46.9174μs | 21.3141 KOps/s | 16.2020 KOps/s | |
| test_compile_indexing[slice-tensordict-compile] | 0.1033ms | 59.0415μs | 16.9372 KOps/s | 16.7354 KOps/s | |
| test_compile_indexing[slice-tensordict-eager] | 0.2487ms | 28.3937μs | 35.2190 KOps/s | 35.7888 KOps/s | |
| test_compile_indexing[slice-tensorclass-compile] | 91.7010μs | 45.7263μs | 21.8693 KOps/s | 21.6439 KOps/s | |
| test_compile_indexing[slice-tensorclass-eager] | 0.2546ms | 22.7397μs | 43.9760 KOps/s | 43.6516 KOps/s | |
| test_compile_indexing[slice-pytree-compile] | 82.1610μs | 45.4322μs | 22.0108 KOps/s | 21.1889 KOps/s | |
| test_compile_indexing[slice-pytree-eager] | 0.2896ms | 23.3441μs | 42.8373 KOps/s | 43.7857 KOps/s | |
| test_compile_indexing[int-tensordict-compile] | 0.1471ms | 56.7202μs | 17.6304 KOps/s | 16.1718 KOps/s | |
| test_compile_indexing[int-tensordict-eager] | 0.3259ms | 27.3837μs | 36.5181 KOps/s | 36.1386 KOps/s | |
| test_compile_indexing[int-tensorclass-compile] | 99.6110μs | 47.3749μs | 21.1082 KOps/s | 21.2361 KOps/s | |
| test_compile_indexing[int-tensorclass-eager] | 0.2647ms | 22.6843μs | 44.0834 KOps/s | 43.6273 KOps/s | |
| test_compile_indexing[int-pytree-compile] | 0.1002ms | 45.4230μs | 22.0153 KOps/s | 20.6393 KOps/s | |
| test_compile_indexing[int-pytree-eager] | 0.2634ms | 23.0848μs | 43.3185 KOps/s | 43.7852 KOps/s | |
| test_mod_add[eager] | 98.9420μs | 51.6466μs | 19.3624 KOps/s | 19.6854 KOps/s | |
| test_mod_add[compile] | 0.1537ms | 0.1085ms | 9.2136 KOps/s | 9.2752 KOps/s | |
| test_mod_add[compile-overhead] | 0.2367ms | 0.1495ms | 6.6901 KOps/s | 6.5222 KOps/s | |
| test_mod_wrap[eager] | 0.3679ms | 0.2946ms | 3.3950 KOps/s | 3.3953 KOps/s | |
| test_mod_wrap[compile] | 0.4141ms | 0.3462ms | 2.8882 KOps/s | 2.8064 KOps/s | |
| test_mod_wrap[compile-overhead] | 7.3676ms | 4.0520ms | 246.7932 Ops/s | 248.1377 Ops/s | |
| test_mod_wrap_and_backward[eager] | 1.5940ms | 1.4932ms | 669.7190 Ops/s | 620.9907 Ops/s | |
| test_mod_wrap_and_backward[compile] | 1.5177ms | 1.4387ms | 695.0627 Ops/s | 635.1017 Ops/s | |
| test_mod_wrap_and_backward[compile-overhead] | 1.2597ms | 0.8895ms | 1.1243 KOps/s | 987.6368 Ops/s | |
| test_seq_add[eager] | 0.2281ms | 0.1561ms | 6.4044 KOps/s | 6.1516 KOps/s | |
| test_seq_add[compile] | 0.1727ms | 0.1178ms | 8.4875 KOps/s | 8.0126 KOps/s | |
| test_seq_add[compile-overhead] | 0.3629ms | 0.1597ms | 6.2635 KOps/s | 6.1270 KOps/s | |
| test_seq_wrap[eager] | 0.6208ms | 0.5405ms | 1.8502 KOps/s | 1.8478 KOps/s | |
| test_seq_wrap[compile] | 0.4179ms | 0.3646ms | 2.7424 KOps/s | 2.5624 KOps/s | |
| test_seq_wrap[compile-overhead] | 0.3359ms | 0.2644ms | 3.7825 KOps/s | 3.6393 KOps/s | |
| test_func_call_runtime[False-eager] | 0.9160ms | 0.8450ms | 1.1835 KOps/s | 1.1294 KOps/s | |
| test_func_call_runtime[False-compile] | 2.1767ms | 0.9120ms | 1.0964 KOps/s | 1.0809 KOps/s | |
| test_func_call_runtime[False-compile-overhead] | 0.5107ms | 0.4608ms | 2.1701 KOps/s | 2.1281 KOps/s | |
| test_func_call_runtime[True-eager] | 1.1565ms | 1.0824ms | 923.8512 Ops/s | 891.1624 Ops/s | |
| test_func_call_runtime[True-compile] | 0.9975ms | 0.9184ms | 1.0889 KOps/s | 1.0547 KOps/s | |
| test_func_call_runtime[True-compile-overhead] | 0.6077ms | 0.4750ms | 2.1052 KOps/s | 2.0573 KOps/s | |
| test_func_call_cm_runtime[False-eager] | 1.0703ms | 0.8825ms | 1.1332 KOps/s | 1.1790 KOps/s | |
| test_func_call_cm_runtime[False-compile] | 1.1222ms | 0.9104ms | 1.0984 KOps/s | 1.0729 KOps/s | |
| test_func_call_cm_runtime[False-compile-overhead] | 0.5700ms | 0.4639ms | 2.1555 KOps/s | 2.1333 KOps/s | |
| test_func_call_cm_runtime[True-eager] | 1.4812ms | 1.2323ms | 811.5149 Ops/s | 798.6816 Ops/s | |
| test_func_call_cm_runtime[True-compile] | 1.0367ms | 0.9546ms | 1.0476 KOps/s | 1.0230 KOps/s | |
| test_func_call_cm_runtime[True-compile-overhead] | 0.5555ms | 0.5062ms | 1.9754 KOps/s | 1.9401 KOps/s | |
| test_vmap_func_call_cm_runtime[eager] | 2.8455ms | 2.3574ms | 424.2036 Ops/s | 417.2925 Ops/s | |
| test_vmap_func_call_cm_runtime[compile] | 1.0342ms | 0.9710ms | 1.0299 KOps/s | 1.0091 KOps/s | |
| test_vmap_func_call_cm_runtime[compile-overhead] | 0.5764ms | 0.5120ms | 1.9531 KOps/s | 1.9127 KOps/s | |
| test_distributed | 2.5648ms | 0.1585ms | 6.3110 KOps/s | 6.4635 KOps/s | |
| test_tdmodule | 0.5367ms | 28.4505μs | 35.1488 KOps/s | 35.5581 KOps/s | |
| test_tdmodule_dispatch | 76.0910μs | 45.5119μs | 21.9723 KOps/s | 21.5537 KOps/s | |
| test_tdseq | 46.6700μs | 26.9832μs | 37.0601 KOps/s | 36.6511 KOps/s | |
| test_tdseq_dispatch | 80.2710μs | 46.9970μs | 21.2780 KOps/s | 20.7663 KOps/s | |
| test_instantiation_functorch | 2.1531ms | 2.0599ms | 485.4651 Ops/s | 478.9384 Ops/s | |
| test_exec_functorch | 0.2256ms | 0.1810ms | 5.5262 KOps/s | 5.5645 KOps/s | |
| test_exec_functional_call | 0.1974ms | 0.1609ms | 6.2133 KOps/s | 6.1571 KOps/s | |
| test_exec_td_decorator | 0.4534ms | 0.2369ms | 4.2213 KOps/s | 4.1880 KOps/s | |
| test_vmap_mlp_speed_decorator[True-True] | 1.0071ms | 0.8186ms | 1.2217 KOps/s | 1.2011 KOps/s | |
| test_vmap_mlp_speed_decorator[True-False] | 1.0184ms | 0.8179ms | 1.2226 KOps/s | 1.2046 KOps/s | |
| test_vmap_mlp_speed_decorator[False-True] | 0.8677ms | 0.7117ms | 1.4051 KOps/s | 1.3878 KOps/s | |
| test_vmap_mlp_speed_decorator[False-False] | 0.9091ms | 0.7112ms | 1.4061 KOps/s | 1.3935 KOps/s | |
| test_vmap_transformer_speed_decorator[True-True] | 21.5032ms | 20.4717ms | 48.8480 Ops/s | 48.3624 Ops/s | |
| test_vmap_transformer_speed_decorator[True-False] | 21.0272ms | 20.4433ms | 48.9158 Ops/s | 48.3853 Ops/s | |
| test_vmap_transformer_speed_decorator[False-True] | 20.3845ms | 20.2487ms | 49.3860 Ops/s | 48.7421 Ops/s | |
| test_vmap_transformer_speed_decorator[False-False] | 20.9857ms | 20.2961ms | 49.2705 Ops/s | 48.7840 Ops/s | |
| test_to_module_speed[True] | 1.5760ms | 1.4899ms | 671.1729 Ops/s | 680.1013 Ops/s | |
| test_to_module_speed[False] | 1.5545ms | 1.4627ms | 683.6779 Ops/s | 689.3720 Ops/s | |
| test_tc_init | 80.8410μs | 45.0203μs | 22.2122 KOps/s | 21.9499 KOps/s | |
| test_tc_init_tensor_only | 35.5400μs | 9.8438μs | 101.5872 KOps/s | 101.0610 KOps/s | |
| test_tc_init_nested | 0.1338ms | 89.8636μs | 11.1280 KOps/s | 11.0633 KOps/s | |
| test_tc_init_many_fields | 43.8910μs | 16.7990μs | 59.5273 KOps/s | 60.7700 KOps/s | |
| test_tc_first_layer_tensor | 18.1000μs | 1.8702μs | 534.6992 KOps/s | 544.1977 KOps/s | |
| test_tc_first_layer_tensor_only | 8.9287μs | 0.7679μs | 1.3023 MOps/s | 1.3329 MOps/s | |
| test_tc_first_layer_tensor_set | 24.1400μs | 3.9549μs | 252.8489 KOps/s | 252.6352 KOps/s | |
| test_tc_first_layer_tensor_only_set | 29.3800μs | 3.1133μs | 321.2027 KOps/s | 318.1045 KOps/s | |
| test_tc_first_layer_nontensor | 25.9300μs | 6.2719μs | 159.4413 KOps/s | 155.2762 KOps/s | |
| test_tc_second_layer_tensor | 24.7510μs | 4.4221μs | 226.1350 KOps/s | 228.5019 KOps/s | |
| test_tc_second_layer_nontensor | 36.9600μs | 8.6238μs | 115.9580 KOps/s | 109.3909 KOps/s | |
| test_unbind | 0.2468s | 14.0324ms | 71.2635 Ops/s | 55.2049 Ops/s | |
| test_full_like | 9.3068ms | 7.3427ms | 136.1905 Ops/s | 228.7819 Ops/s | |
| test_zeros_like | 4.9627ms | 4.3606ms | 229.3280 Ops/s | 229.2154 Ops/s | |
| test_ones_like | 4.4964ms | 4.3612ms | 229.2930 Ops/s | 229.0003 Ops/s | |
| test_clone | 6.5903ms | 6.4022ms | 156.1964 Ops/s | 155.8008 Ops/s | |
| test_squeeze | 0.1390ms | 14.4407μs | 69.2486 KOps/s | 69.0852 KOps/s | |
| test_unsqueeze | 0.2523ms | 0.1111ms | 8.9993 KOps/s | 9.0204 KOps/s | |
| test_split | 0.2312ms | 0.1845ms | 5.4196 KOps/s | 5.4081 KOps/s | |
| test_permute | 0.2623ms | 0.2019ms | 4.9540 KOps/s | 4.8536 KOps/s | |
| test_stack | 51.5429ms | 51.2394ms | 19.5162 Ops/s | 19.5229 Ops/s | |
| test_cat | 51.9415ms | 51.2992ms | 19.4935 Ops/s | 19.5752 Ops/s | |
| test_sequential_tensordict | 0.2698ms | 0.2140ms | 4.6731 KOps/s | 4.3964 KOps/s | |
| test_sequential_graph_module | 0.5481ms | 0.1189ms | 8.4083 KOps/s | 7.8072 KOps/s | |
| test_nested_tensordict | 0.3400ms | 0.2821ms | 3.5442 KOps/s | 3.4465 KOps/s | |
| test_nested_graph_module | 0.5237ms | 0.1274ms | 7.8487 KOps/s | 7.3580 KOps/s |
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
PR Title Label ErrorUnknown or invalid prefix Current title: Supported PrefixesYour PR title must start with exactly one of these prefixes:
Note: Common variations like singular/plural are supported (e.g., |
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
Stack from ghstack (oldest at bottom):
Send the fully qualified tensorclass type string (e.g.
tensordict.testing.MyDistData) as metadata during init_remote andbroadcast. On the receiver side, from_remote_init and broadcast resolve
the type and wrap the reconstructed TensorDict into the original
tensorclass.
_tensorclass_typeprivate param toinit_remote/broadcastin TensorDictBase, injected into consolidation metadata.
_resolve_tensorclass_typehelper to import a class from itsdotted path.
init_remote,broadcast,from_remote_initontensorclass to inject / recover the type automatically.
tensordict/testing.pywithMyDistDataso both sides ofa multiprocessing test can import the same class.
type(result) is MyDistData.Co-authored-by: Cursor cursoragent@cursor.com