Skip to content

[Compile] Make lock_() and other _as_context_manager ops compile-safe - #1714

Merged
vmoens merged 2 commits into
gh/vmoens/94/basefrom
gh/vmoens/94/head
May 24, 2026
Merged

[Compile] Make lock_() and other _as_context_manager ops compile-safe#1714
vmoens merged 2 commits into
gh/vmoens/94/basefrom
gh/vmoens/94/head

Conversation

@vmoens

@vmoens vmoens commented May 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

The _as_context_manager decorator (used by lock_/unlock_,
to_module, transpose, flatten, etc.) wrote a
weakref.ref(self) into _last_op so the corresponding __exit__
could call the inverse op. Under torch.compile that weakref becomes
a WeakRefVariable that Dynamo cannot reconstruct when the holding
TD flows out of the compiled region, which forced callers to defensively
call clear_refs_for_compile_() (and previously emitted a strong
warning around lock_ itself).

Use a strong-ref closure (_strong_ref(obj) returning lambda: obj)
in place of weakref.ref when is_compiling(). The closure is a
plain Python callable Dynamo can handle, and it's safe because the
compiled graph is transient: the closure is dropped as soon as the
outer call ends, so no real lifecycle cycle can form.

Also:

  • Add _last_op = None as a class-level default on TensorDictBase
    so __enter__ can always read it (the compile path no longer writes
    it on every decorated call).
  • Update the user-facing lock_ warning to describe the real
    remaining caveat under compile: the parent/child unlock invariant is
    not enforced, the previous weakref-reconstruction failure is gone.

Adds a regression test asserting that a TD locked inside a compiled
function does not leak a real weakref.ref into _last_op, and
that the strong-ref closure resolves back to the same TD.

Authored with Claude.

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

@vmoens
Update
6927029 
[ghstack-poisoned]
@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 May 24, 2026
@github-actions github-actions Bot added Compile torch.compile related Test and removed CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. labels May 24, 2026
This was referenced May 24, 2026
Merged
Merged
@github-actions

github-actions Bot commented May 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 341. Improved: $\large\color{#35bf28}14$. Worsened: $\large\color{#d91a1a}19$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_collector_preallocated_write[stack_out-cpu-32-2] 0.2041ms 0.1707ms 5.8593 KOps/s 5.9552 KOps/s $\color{#d91a1a}-1.61\%$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.3984ms 0.2975ms 3.3618 KOps/s 3.3808 KOps/s $\color{#d91a1a}-0.56\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8250ms 0.7764ms 1.2880 KOps/s 1.2834 KOps/s $\color{#35bf28}+0.36\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.4535ms 0.3422ms 2.9222 KOps/s 2.9208 KOps/s $\color{#35bf28}+0.05\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.2068ms 1.0954ms 912.9242 Ops/s 900.4891 Ops/s $\color{#35bf28}+1.38\%$
test_collector_preallocated_write[stack_out-cpu-4096-32] 7.5471ms 7.3866ms 135.3806 Ops/s 133.6528 Ops/s $\color{#35bf28}+1.29\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.3068ms 0.2546ms 3.9276 KOps/s 3.9870 KOps/s $\color{#d91a1a}-1.49\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.5153ms 0.3753ms 2.6642 KOps/s 2.6509 KOps/s $\color{#35bf28}+0.50\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.9936ms 0.8272ms 1.2089 KOps/s 1.2134 KOps/s $\color{#d91a1a}-0.37\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.6722ms 0.2529ms 3.9546 KOps/s 3.8575 KOps/s $\color{#35bf28}+2.52\%$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.8155ms 0.3781ms 2.6448 KOps/s 2.6169 KOps/s $\color{#35bf28}+1.06\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 1.3164ms 0.8857ms 1.1290 KOps/s 1.1256 KOps/s $\color{#35bf28}+0.30\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.2658ms 0.2208ms 4.5292 KOps/s 4.2456 KOps/s $\textbf{\color{#35bf28}+6.68\%}$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 1.2902ms 0.8500ms 1.1765 KOps/s 1.1212 KOps/s $\color{#35bf28}+4.93\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.8229ms 3.3903ms 294.9576 Ops/s 285.5210 Ops/s $\color{#35bf28}+3.31\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.8024ms 0.3727ms 2.6835 KOps/s 2.5571 KOps/s $\color{#35bf28}+4.94\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 2.0200ms 1.7734ms 563.8883 Ops/s 564.7971 Ops/s $\color{#d91a1a}-0.16\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 10.7703ms 10.4543ms 95.6545 Ops/s 96.8198 Ops/s $\color{#d91a1a}-1.20\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.5090ms 0.4063ms 2.4615 KOps/s 2.4437 KOps/s $\color{#35bf28}+0.73\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.6717ms 1.5151ms 660.0148 Ops/s 668.3105 Ops/s $\color{#d91a1a}-1.24\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 6.1140ms 5.9651ms 167.6405 Ops/s 172.4337 Ops/s $\color{#d91a1a}-2.78\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.8258ms 0.4038ms 2.4762 KOps/s 2.5308 KOps/s $\color{#d91a1a}-2.16\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.9694ms 1.5216ms 657.2176 Ops/s 665.4453 Ops/s $\color{#d91a1a}-1.24\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 6.4010ms 5.9804ms 167.2130 Ops/s 171.9589 Ops/s $\color{#d91a1a}-2.76\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.2243ms 0.1767ms 5.6602 KOps/s 5.7482 KOps/s $\color{#d91a1a}-1.53\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 1.1150ms 0.6894ms 1.4505 KOps/s 1.4826 KOps/s $\color{#d91a1a}-2.16\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 3.1721ms 2.7224ms 367.3266 Ops/s 375.3490 Ops/s $\color{#d91a1a}-2.14\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.7577ms 0.3213ms 3.1122 KOps/s 3.0961 KOps/s $\color{#35bf28}+0.52\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 1.8295ms 1.5947ms 627.0588 Ops/s 689.1588 Ops/s $\textbf{\color{#d91a1a}-9.01\%}$
test_collector_preallocated_write[copy_at-cpu-4096-32] 9.8288ms 9.5886ms 104.2902 Ops/s 104.3709 Ops/s $\color{#d91a1a}-0.08\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.4228ms 0.3334ms 2.9992 KOps/s 3.0615 KOps/s $\color{#d91a1a}-2.03\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.4021ms 1.2455ms 802.9101 Ops/s 820.1068 Ops/s $\color{#d91a1a}-2.10\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 5.0289ms 4.8817ms 204.8463 Ops/s 211.2288 Ops/s $\color{#d91a1a}-3.02\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.4224ms 0.3342ms 2.9923 KOps/s 3.0549 KOps/s $\color{#d91a1a}-2.05\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.3626ms 1.2503ms 799.7944 Ops/s 824.7960 Ops/s $\color{#d91a1a}-3.03\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 5.0588ms 4.9094ms 203.6928 Ops/s 210.0354 Ops/s $\color{#d91a1a}-3.02\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1898ms 0.1253ms 7.9786 KOps/s 8.2817 KOps/s $\color{#d91a1a}-3.66\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5874ms 0.4932ms 2.0276 KOps/s 2.1114 KOps/s $\color{#d91a1a}-3.97\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 2.0836ms 1.9663ms 508.5693 Ops/s 532.8539 Ops/s $\color{#d91a1a}-4.56\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.3368ms 0.2567ms 3.8962 KOps/s 3.7948 KOps/s $\color{#35bf28}+2.67\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.4478ms 1.2378ms 807.8930 Ops/s 825.8883 Ops/s $\color{#d91a1a}-2.18\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 8.8870ms 8.7016ms 114.9208 Ops/s 113.6227 Ops/s $\color{#35bf28}+1.14\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.3744ms 0.2896ms 3.4529 KOps/s 3.5515 KOps/s $\color{#d91a1a}-2.77\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.1785ms 1.0907ms 916.8148 Ops/s 960.2758 Ops/s $\color{#d91a1a}-4.53\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.4449ms 4.2238ms 236.7560 Ops/s 251.1896 Ops/s $\textbf{\color{#d91a1a}-5.75\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.4018ms 0.2902ms 3.4462 KOps/s 3.6081 KOps/s $\color{#d91a1a}-4.49\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.2431ms 1.0890ms 918.3025 Ops/s 973.5676 Ops/s $\textbf{\color{#d91a1a}-5.68\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.3605ms 4.2250ms 236.6857 Ops/s 249.6316 Ops/s $\textbf{\color{#d91a1a}-5.19\%}$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 16.4190ms 16.1293ms 61.9989 Ops/s 61.6491 Ops/s $\color{#35bf28}+0.57\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 32.9381ms 32.6863ms 30.5938 Ops/s 30.5987 Ops/s $\color{#d91a1a}-0.02\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 86.5686ms 86.2660ms 11.5921 Ops/s 11.7461 Ops/s $\color{#d91a1a}-1.31\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2044s 0.2034s 4.9173 Ops/s 4.9622 Ops/s $\color{#d91a1a}-0.90\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 14.1814ms 13.7194ms 72.8897 Ops/s 72.8352 Ops/s $\color{#35bf28}+0.07\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 27.2406ms 26.8301ms 37.2715 Ops/s 37.1591 Ops/s $\color{#35bf28}+0.30\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 15.2603ms 13.8882ms 72.0036 Ops/s 71.7692 Ops/s $\color{#35bf28}+0.33\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 27.6296ms 27.1105ms 36.8860 Ops/s 36.5873 Ops/s $\color{#35bf28}+0.82\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 53.3350ms 52.9591ms 18.8825 Ops/s 18.9589 Ops/s $\color{#d91a1a}-0.40\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1084s 0.1078s 9.2777 Ops/s 9.4475 Ops/s $\color{#d91a1a}-1.80\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1494s 0.1487s 6.7263 Ops/s 6.8111 Ops/s $\color{#d91a1a}-1.25\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3287s 0.3271s 3.0570 Ops/s 3.0937 Ops/s $\color{#d91a1a}-1.18\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 80.0599ms 78.8439ms 12.6833 Ops/s 13.5367 Ops/s $\textbf{\color{#d91a1a}-6.30\%}$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1587s 0.1562s 6.4038 Ops/s 6.6364 Ops/s $\color{#d91a1a}-3.50\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 80.2228ms 77.8568ms 12.8441 Ops/s 13.2814 Ops/s $\color{#d91a1a}-3.29\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1614s 0.1554s 6.4331 Ops/s 6.5388 Ops/s $\color{#d91a1a}-1.62\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.9058ms 30.5175ms 32.7680 Ops/s 32.9558 Ops/s $\color{#d91a1a}-0.57\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 62.4694ms 62.1493ms 16.0903 Ops/s 16.3523 Ops/s $\color{#d91a1a}-1.60\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1282s 0.1218s 8.2127 Ops/s 8.3765 Ops/s $\color{#d91a1a}-1.96\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2719s 0.2699s 3.7046 Ops/s 3.6169 Ops/s $\color{#35bf28}+2.43\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 44.8732ms 44.5591ms 22.4421 Ops/s 23.1592 Ops/s $\color{#d91a1a}-3.10\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 89.7250ms 89.3308ms 11.1943 Ops/s 11.5834 Ops/s $\color{#d91a1a}-3.36\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 45.4116ms 44.9065ms 22.2685 Ops/s 23.0468 Ops/s $\color{#d91a1a}-3.38\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 90.4785ms 90.0267ms 11.1078 Ops/s 11.4770 Ops/s $\color{#d91a1a}-3.22\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 21.5027ms 21.1159ms 47.3577 Ops/s 48.8618 Ops/s $\color{#d91a1a}-3.08\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 43.6489ms 43.2940ms 23.0979 Ops/s 23.8917 Ops/s $\color{#d91a1a}-3.32\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1132s 0.1115s 8.9702 Ops/s 9.1687 Ops/s $\color{#d91a1a}-2.17\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2509s 0.2498s 4.0039 Ops/s 3.9256 Ops/s $\color{#35bf28}+1.99\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 39.3005ms 38.7959ms 25.7759 Ops/s 27.1185 Ops/s $\color{#d91a1a}-4.95\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 78.2619ms 77.6683ms 12.8753 Ops/s 13.5486 Ops/s $\color{#d91a1a}-4.97\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 39.3788ms 38.8033ms 25.7710 Ops/s 27.0999 Ops/s $\color{#d91a1a}-4.90\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 78.3882ms 77.7795ms 12.8569 Ops/s 13.5684 Ops/s $\textbf{\color{#d91a1a}-5.24\%}$
test_plain_set_nested 46.8010μs 15.5043μs 64.4984 KOps/s 65.1685 KOps/s $\color{#d91a1a}-1.03\%$
test_plain_set_stack_nested 39.2010μs 15.2806μs 65.4426 KOps/s 66.3581 KOps/s $\color{#d91a1a}-1.38\%$
test_plain_set_nested_inplace 0.4321ms 16.9007μs 59.1692 KOps/s 59.6107 KOps/s $\color{#d91a1a}-0.74\%$
test_plain_set_stack_nested_inplace 0.4544ms 16.5124μs 60.5607 KOps/s 60.6157 KOps/s $\color{#d91a1a}-0.09\%$
test_items 36.2300μs 5.9477μs 168.1322 KOps/s 165.6897 KOps/s $\color{#35bf28}+1.47\%$
test_items_nested 0.8859ms 0.4634ms 2.1578 KOps/s 2.1770 KOps/s $\color{#d91a1a}-0.88\%$
test_items_nested_locked 0.8868ms 0.4673ms 2.1401 KOps/s 2.1516 KOps/s $\color{#d91a1a}-0.53\%$
test_items_nested_leaf 0.5146ms 96.4760μs 10.3653 KOps/s 10.4290 KOps/s $\color{#d91a1a}-0.61\%$
test_items_stack_nested 0.8835ms 0.4642ms 2.1541 KOps/s 2.1646 KOps/s $\color{#d91a1a}-0.48\%$
test_items_stack_nested_leaf 0.5161ms 97.4379μs 10.2630 KOps/s 10.4596 KOps/s $\color{#d91a1a}-1.88\%$
test_items_stack_nested_locked 0.8853ms 0.4702ms 2.1268 KOps/s 2.1518 KOps/s $\color{#d91a1a}-1.16\%$
test_keys 44.1510μs 4.2362μs 236.0602 KOps/s 237.9018 KOps/s $\color{#d91a1a}-0.77\%$
test_keys_nested 0.5570ms 0.1306ms 7.6562 KOps/s 7.7128 KOps/s $\color{#d91a1a}-0.73\%$
test_keys_nested_locked 0.7573ms 0.1393ms 7.1795 KOps/s 7.2693 KOps/s $\color{#d91a1a}-1.23\%$
test_keys_nested_leaf 0.5363ms 0.1206ms 8.2923 KOps/s 8.2813 KOps/s $\color{#35bf28}+0.13\%$
test_keys_stack_nested 0.5582ms 0.1305ms 7.6623 KOps/s 7.7090 KOps/s $\color{#d91a1a}-0.61\%$
test_keys_stack_nested_leaf 0.5416ms 0.1203ms 8.3131 KOps/s 8.3694 KOps/s $\color{#d91a1a}-0.67\%$
test_keys_stack_nested_locked 0.5756ms 0.1388ms 7.2034 KOps/s 7.2673 KOps/s $\color{#d91a1a}-0.88\%$
test_values 85.4036μs 1.0269μs 973.7959 KOps/s 988.7598 KOps/s $\color{#d91a1a}-1.51\%$
test_values_nested 0.4735ms 52.7846μs 18.9449 KOps/s 19.1385 KOps/s $\color{#d91a1a}-1.01\%$
test_values_nested_locked 0.4729ms 55.8610μs 17.9016 KOps/s 18.0396 KOps/s $\color{#d91a1a}-0.77\%$
test_values_nested_leaf 94.1010μs 60.3317μs 16.5750 KOps/s 16.7750 KOps/s $\color{#d91a1a}-1.19\%$
test_values_stack_nested 0.4833ms 52.6693μs 18.9864 KOps/s 19.1526 KOps/s $\color{#d91a1a}-0.87\%$
test_values_stack_nested_leaf 0.4922ms 60.0958μs 16.6401 KOps/s 16.7928 KOps/s $\color{#d91a1a}-0.91\%$
test_values_stack_nested_locked 0.4836ms 55.4827μs 18.0236 KOps/s 18.0480 KOps/s $\color{#d91a1a}-0.14\%$
test_membership 71.5995μs 0.8326μs 1.2010 MOps/s 1.1969 MOps/s $\color{#35bf28}+0.35\%$
test_membership_nested 0.4276ms 2.9733μs 336.3305 KOps/s 336.6685 KOps/s $\color{#d91a1a}-0.10\%$
test_membership_nested_leaf 37.6400μs 2.9656μs 337.1996 KOps/s 335.9212 KOps/s $\color{#35bf28}+0.38\%$
test_membership_stacked_nested 0.4457ms 2.8642μs 349.1371 KOps/s 345.3018 KOps/s $\color{#35bf28}+1.11\%$
test_membership_stacked_nested_leaf 38.3300μs 2.8576μs 349.9467 KOps/s 349.6097 KOps/s $\color{#35bf28}+0.10\%$
test_membership_nested_last 0.4308ms 4.4022μs 227.1610 KOps/s 226.3524 KOps/s $\color{#35bf28}+0.36\%$
test_membership_nested_leaf_last 31.9700μs 4.4149μs 226.5079 KOps/s 226.8635 KOps/s $\color{#d91a1a}-0.16\%$
test_membership_stacked_nested_last 41.0500μs 4.3377μs 230.5364 KOps/s 233.6904 KOps/s $\color{#d91a1a}-1.35\%$
test_membership_stacked_nested_leaf_last 43.1110μs 4.2984μs 232.6463 KOps/s 234.1453 KOps/s $\color{#d91a1a}-0.64\%$
test_nested_getleaf 0.4488ms 21.1333μs 47.3186 KOps/s 45.8364 KOps/s $\color{#35bf28}+3.23\%$
test_nested_get 0.4521ms 20.0618μs 49.8461 KOps/s 49.1525 KOps/s $\color{#35bf28}+1.41\%$
test_stacked_getleaf 0.4475ms 21.0373μs 47.5345 KOps/s 46.1847 KOps/s $\color{#35bf28}+2.92\%$
test_stacked_get 51.9710μs 19.9760μs 50.0600 KOps/s 49.2196 KOps/s $\color{#35bf28}+1.71\%$
test_nested_getitemleaf 0.4534ms 21.8217μs 45.8260 KOps/s 45.0847 KOps/s $\color{#35bf28}+1.64\%$
test_nested_getitem 0.4445ms 20.4050μs 49.0075 KOps/s 47.3403 KOps/s $\color{#35bf28}+3.52\%$
test_stacked_getitemleaf 0.4459ms 21.4212μs 46.6827 KOps/s 45.1571 KOps/s $\color{#35bf28}+3.38\%$
test_stacked_getitem 0.4505ms 20.4159μs 48.9815 KOps/s 47.6517 KOps/s $\color{#35bf28}+2.79\%$
test_lock_nested 0.9209ms 0.4808ms 2.0801 KOps/s 2.0932 KOps/s $\color{#d91a1a}-0.63\%$
test_lock_stack_nested 0.9106ms 0.4860ms 2.0577 KOps/s 2.0972 KOps/s $\color{#d91a1a}-1.89\%$
test_unlock_nested 0.4605ms 0.3911ms 2.5567 KOps/s 2.6095 KOps/s $\color{#d91a1a}-2.02\%$
test_unlock_stack_nested 0.4671ms 0.3937ms 2.5398 KOps/s 2.5905 KOps/s $\color{#d91a1a}-1.96\%$
test_flatten_speed 0.1582ms 0.1208ms 8.2786 KOps/s 8.2026 KOps/s $\color{#35bf28}+0.93\%$
test_unflatten_speed 0.7766ms 0.5700ms 1.7543 KOps/s 1.7499 KOps/s $\color{#35bf28}+0.25\%$
test_common_ops 99.9456ms 1.7978ms 556.2327 Ops/s 238.2231 Ops/s $\textbf{\color{#35bf28}+133.49\%}$
test_creation 0.1086ms 3.1316μs 319.3218 KOps/s 324.6816 KOps/s $\color{#d91a1a}-1.65\%$
test_creation_empty 43.4310μs 6.9753μs 143.3634 KOps/s 143.7426 KOps/s $\color{#d91a1a}-0.26\%$
test_creation_nested_1 73.4710μs 11.1069μs 90.0341 KOps/s 87.2439 KOps/s $\color{#35bf28}+3.20\%$
test_creation_nested_2 33.8810μs 13.5969μs 73.5464 KOps/s 74.6020 KOps/s $\color{#d91a1a}-1.41\%$
test_creation_many_keys[10] 45.9910μs 20.7296μs 48.2402 KOps/s 47.3943 KOps/s $\color{#35bf28}+1.78\%$
test_creation_many_keys[50] 0.1219ms 90.6614μs 11.0301 KOps/s 11.0711 KOps/s $\color{#d91a1a}-0.37\%$
test_creation_many_keys[100] 0.2126ms 0.1776ms 5.6319 KOps/s 5.6441 KOps/s $\color{#d91a1a}-0.22\%$
test_creation_nested_many_keys[10] 74.1010μs 45.1867μs 22.1304 KOps/s 21.9235 KOps/s $\color{#35bf28}+0.94\%$
test_creation_nested_many_keys[50] 0.2382ms 0.1834ms 5.4516 KOps/s 5.3979 KOps/s $\color{#35bf28}+1.00\%$
test_clone 42.0510μs 13.6350μs 73.3405 KOps/s 74.8677 KOps/s $\color{#d91a1a}-2.04\%$
test_getitem[int] 1.4879ms 15.5561μs 64.2833 KOps/s 65.6926 KOps/s $\color{#d91a1a}-2.15\%$
test_getitem[slice_int] 0.1410ms 24.8318μs 40.2710 KOps/s 41.0052 KOps/s $\color{#d91a1a}-1.79\%$
test_getitem[range] 0.1824ms 65.0313μs 15.3772 KOps/s 15.1989 KOps/s $\color{#35bf28}+1.17\%$
test_getitem[tuple] 0.1466ms 24.8967μs 40.1660 KOps/s 41.1177 KOps/s $\color{#d91a1a}-2.31\%$
test_getitem[list] 0.1881ms 59.3662μs 16.8446 KOps/s 15.8810 KOps/s $\textbf{\color{#35bf28}+6.07\%}$
test_setitem_dim[int] 45.6710μs 26.5638μs 37.6452 KOps/s 36.5266 KOps/s $\color{#35bf28}+3.06\%$
test_setitem_dim[slice_int] 65.1220μs 43.2025μs 23.1468 KOps/s 22.0765 KOps/s $\color{#35bf28}+4.85\%$
test_setitem_dim[range] 0.1371ms 96.8169μs 10.3288 KOps/s 9.7270 KOps/s $\textbf{\color{#35bf28}+6.19\%}$
test_setitem_dim[tuple] 69.3110μs 39.8180μs 25.1143 KOps/s 25.3154 KOps/s $\color{#d91a1a}-0.79\%$
test_setitem 52.7210μs 17.8334μs 56.0746 KOps/s 53.6384 KOps/s $\color{#35bf28}+4.54\%$
test_set 51.5210μs 17.2290μs 58.0417 KOps/s 56.1196 KOps/s $\color{#35bf28}+3.43\%$
test_set_shared 0.4892ms 0.2054ms 4.8692 KOps/s 4.8014 KOps/s $\color{#35bf28}+1.41\%$
test_update 0.1991ms 22.3238μs 44.7952 KOps/s 45.0700 KOps/s $\color{#d91a1a}-0.61\%$
test_update_nested 72.2410μs 33.6769μs 29.6940 KOps/s 29.7414 KOps/s $\color{#d91a1a}-0.16\%$
test_update__nested 0.4804ms 33.3487μs 29.9862 KOps/s 29.8975 KOps/s $\color{#35bf28}+0.30\%$
test_set_nested 55.4810μs 19.1041μs 52.3447 KOps/s 52.9548 KOps/s $\color{#d91a1a}-1.15\%$
test_set_nested_new 49.8210μs 24.1570μs 41.3958 KOps/s 42.6057 KOps/s $\color{#d91a1a}-2.84\%$
test_select 93.5810μs 40.9300μs 24.4320 KOps/s 24.8990 KOps/s $\color{#d91a1a}-1.88\%$
test_select_nested 0.1157ms 74.8287μs 13.3638 KOps/s 13.5123 KOps/s $\color{#d91a1a}-1.10\%$
test_exclude_nested 0.1332ms 91.6303μs 10.9134 KOps/s 10.9989 KOps/s $\color{#d91a1a}-0.78\%$
test_empty[True] 0.4671ms 0.4020ms 2.4876 KOps/s 2.5188 KOps/s $\color{#d91a1a}-1.24\%$
test_empty[False] 11.9267μs 1.3987μs 714.9712 KOps/s 715.2401 KOps/s $\color{#d91a1a}-0.04\%$
test_to 0.1142ms 70.9107μs 14.1023 KOps/s 14.4258 KOps/s $\color{#d91a1a}-2.24\%$
test_to_nonblocking 0.1204ms 69.1794μs 14.4552 KOps/s 14.4326 KOps/s $\color{#35bf28}+0.16\%$
test_unbind_speed 0.4314ms 0.3430ms 2.9154 KOps/s 2.9881 KOps/s $\color{#d91a1a}-2.43\%$
test_unbind_speed_stack0 0.3819ms 0.3383ms 2.9563 KOps/s 3.0122 KOps/s $\color{#d91a1a}-1.85\%$
test_unbind_speed_stack1 0.1145s 0.8650ms 1.1560 KOps/s 1.1606 KOps/s $\color{#d91a1a}-0.40\%$
test_split 1.2651ms 1.1434ms 874.6110 Ops/s 873.1088 Ops/s $\color{#35bf28}+0.17\%$
test_chunk 0.1143s 1.2307ms 812.5451 Ops/s 666.5483 Ops/s $\textbf{\color{#35bf28}+21.90\%}$
test_to_cpu_blocking 20.3776ms 20.0230ms 49.9426 Ops/s 49.8701 Ops/s $\color{#35bf28}+0.15\%$
test_to_cpu_global_sync 12.1817ms 12.0100ms 83.2638 Ops/s 84.8509 Ops/s $\color{#d91a1a}-1.87\%$
test_to_cpu_event_sync 13.5456ms 12.9989ms 76.9298 Ops/s 78.7156 Ops/s $\color{#d91a1a}-2.27\%$
test_to_cpu_default 13.3161ms 12.9321ms 77.3273 Ops/s 78.4369 Ops/s $\color{#d91a1a}-1.41\%$
test_consolidate[False-None] 0.1440s 4.7652ms 209.8549 Ops/s 240.7869 Ops/s $\textbf{\color{#d91a1a}-12.85\%}$
test_consolidate[default-None] 3.7520ms 2.1847ms 457.7253 Ops/s 489.1509 Ops/s $\textbf{\color{#d91a1a}-6.42\%}$
test_consolidate[reduce-overhead-None] 3.5695ms 2.1240ms 470.8062 Ops/s 505.5979 Ops/s $\textbf{\color{#d91a1a}-6.88\%}$
test_consolidate_njt[False-None] 0.2079s 10.3289ms 96.8156 Ops/s 113.1138 Ops/s $\textbf{\color{#d91a1a}-14.41\%}$
test_to[False-False-None] 2.6236ms 2.2086ms 452.7695 Ops/s 455.9428 Ops/s $\color{#d91a1a}-0.70\%$
test_to[True-False-None] 2.3758ms 1.9330ms 517.3358 Ops/s 535.1348 Ops/s $\color{#d91a1a}-3.33\%$
test_to[within-False-None] 6.6893ms 6.1995ms 161.3034 Ops/s 162.1673 Ops/s $\color{#d91a1a}-0.53\%$
test_to[True-default-None] 10.9976ms 10.7748ms 92.8096 Ops/s 91.0877 Ops/s $\color{#35bf28}+1.89\%$
test_to_njt[False-False-None] 8.9743ms 8.6448ms 115.6764 Ops/s 114.2946 Ops/s $\color{#35bf28}+1.21\%$
test_to_njt[True-False-None] 7.5465ms 7.1238ms 140.3749 Ops/s 141.9333 Ops/s $\color{#d91a1a}-1.10\%$
test_to_njt[within-False-None] 15.9284ms 15.6889ms 63.7392 Ops/s 60.8131 Ops/s $\color{#35bf28}+4.81\%$
test_creation[device0] 0.5744ms 0.1192ms 8.3869 KOps/s 8.3260 KOps/s $\color{#35bf28}+0.73\%$
test_creation_from_tensor 0.4334ms 0.1185ms 8.4373 KOps/s 8.8020 KOps/s $\color{#d91a1a}-4.14\%$
test_add_one[memmap_tensor0] 0.2448ms 7.0177μs 142.4961 KOps/s 141.6393 KOps/s $\color{#35bf28}+0.60\%$
test_contiguous[memmap_tensor0] 14.6800μs 0.6593μs 1.5168 MOps/s 2.2116 MOps/s $\textbf{\color{#d91a1a}-31.41\%}$
test_stack[memmap_tensor0] 27.5000μs 4.4752μs 223.4517 KOps/s 237.7649 KOps/s $\textbf{\color{#d91a1a}-6.02\%}$
test_memmaptd_index 1.2716ms 0.2864ms 3.4914 KOps/s 3.7060 KOps/s $\textbf{\color{#d91a1a}-5.79\%}$
test_memmaptd_index_astensor 0.8023ms 0.3910ms 2.5576 KOps/s 2.6577 KOps/s $\color{#d91a1a}-3.77\%$
test_memmaptd_index_op 1.0558ms 0.6576ms 1.5206 KOps/s 1.5417 KOps/s $\color{#d91a1a}-1.37\%$
test_serialize_model 0.3233s 0.1619s 6.1781 Ops/s 7.4333 Ops/s $\textbf{\color{#d91a1a}-16.89\%}$
test_serialize_model_pickle 1.3491s 1.2102s 0.8263 Ops/s 0.8267 Ops/s $\color{#d91a1a}-0.05\%$
test_serialize_weights 0.1376s 0.1338s 7.4726 Ops/s 7.4871 Ops/s $\color{#d91a1a}-0.19\%$
test_serialize_weights_returnearly 0.4189s 84.0617ms 11.8960 Ops/s 15.4387 Ops/s $\textbf{\color{#d91a1a}-22.95\%}$
test_serialize_weights_pickle 1.3844s 1.2191s 0.8203 Ops/s 0.8237 Ops/s $\color{#d91a1a}-0.41\%$
test_reshape_pytree 0.2119ms 34.1424μs 29.2891 KOps/s 30.5979 KOps/s $\color{#d91a1a}-4.28\%$
test_reshape_td 0.1975ms 45.4994μs 21.9783 KOps/s 22.3315 KOps/s $\color{#d91a1a}-1.58\%$
test_view_pytree 0.2141ms 33.1661μs 30.1513 KOps/s 30.9605 KOps/s $\color{#d91a1a}-2.61\%$
test_view_td 80.0320μs 54.7574μs 18.2624 KOps/s 18.9779 KOps/s $\color{#d91a1a}-3.77\%$
test_unbind_pytree 0.2326ms 37.7071μs 26.5202 KOps/s 26.8387 KOps/s $\color{#d91a1a}-1.19\%$
test_unbind_td 96.2220μs 50.6105μs 19.7588 KOps/s 19.9854 KOps/s $\color{#d91a1a}-1.13\%$
test_split_pytree 0.2458ms 43.5349μs 22.9701 KOps/s 23.5426 KOps/s $\color{#d91a1a}-2.43\%$
test_split_td 0.1419ms 66.1696μs 15.1127 KOps/s 15.8734 KOps/s $\color{#d91a1a}-4.79\%$
test_add_pytree 0.1933ms 44.2535μs 22.5971 KOps/s 22.8437 KOps/s $\color{#d91a1a}-1.08\%$
test_add_td 0.1465ms 56.0094μs 17.8541 KOps/s 17.5933 KOps/s $\color{#35bf28}+1.48\%$
test_compile_add_one_nested[tensordict-compile] 0.3535ms 0.2307ms 4.3340 KOps/s 4.2319 KOps/s $\color{#35bf28}+2.41\%$
test_compile_add_one_nested[tensordict-eager] 0.3340ms 0.2028ms 4.9308 KOps/s 4.9576 KOps/s $\color{#d91a1a}-0.54\%$
test_compile_add_one_nested[pytree-compile] 0.3109ms 0.1353ms 7.3901 KOps/s 7.1313 KOps/s $\color{#35bf28}+3.63\%$
test_compile_add_one_nested[pytree-eager] 0.4402ms 0.1831ms 5.4620 KOps/s 5.5220 KOps/s $\color{#d91a1a}-1.09\%$
test_compile_copy_nested[tensordict-compile] 0.1487ms 14.6915μs 68.0668 KOps/s 67.0951 KOps/s $\color{#35bf28}+1.45\%$
test_compile_copy_nested[tensordict-eager] 0.2685ms 53.4649μs 18.7039 KOps/s 18.8178 KOps/s $\color{#d91a1a}-0.61\%$
test_compile_copy_nested[pytree-compile] 0.1198ms 15.8709μs 63.0084 KOps/s 63.5421 KOps/s $\color{#d91a1a}-0.84\%$
test_compile_copy_nested[pytree-eager] 0.3795ms 68.9890μs 14.4951 KOps/s 14.5693 KOps/s $\color{#d91a1a}-0.51\%$
test_compile_add_one_flat[tensordict-compile] 0.3982ms 0.2486ms 4.0218 KOps/s 3.8642 KOps/s $\color{#35bf28}+4.08\%$
test_compile_add_one_flat[tensordict-eager] 0.3510ms 0.2820ms 3.5455 KOps/s 3.4605 KOps/s $\color{#35bf28}+2.46\%$
test_compile_add_one_flat[tensorclass-compile] 0.3405ms 0.1832ms 5.4571 KOps/s 5.2185 KOps/s $\color{#35bf28}+4.57\%$
test_compile_add_one_flat[tensorclass-eager] 0.1160ms 74.9083μs 13.3497 KOps/s 13.4490 KOps/s $\color{#d91a1a}-0.74\%$
test_compile_add_one_flat[pytree-compile] 0.4155ms 0.1914ms 5.2260 KOps/s 5.0143 KOps/s $\color{#35bf28}+4.22\%$
test_compile_add_one_flat[pytree-eager] 0.8023ms 0.5377ms 1.8599 KOps/s 1.9250 KOps/s $\color{#d91a1a}-3.38\%$
test_compile_add_self_flat[tensordict-eager] 0.4665ms 0.3297ms 3.0327 KOps/s 1.3309 KOps/s $\textbf{\color{#35bf28}+127.88\%}$
test_compile_add_self_flat[tensordict-compile] 0.3111ms 0.2482ms 4.0287 KOps/s 3.6907 KOps/s $\textbf{\color{#35bf28}+9.16\%}$
test_compile_add_self_flat[tensorclass-eager] 0.1293ms 91.9172μs 10.8794 KOps/s 11.0019 KOps/s $\color{#d91a1a}-1.11\%$
test_compile_add_self_flat[tensorclass-compile] 0.3229ms 0.1860ms 5.3776 KOps/s 5.0137 KOps/s $\textbf{\color{#35bf28}+7.26\%}$
test_compile_add_self_flat[pytree-eager] 0.6615ms 0.4419ms 2.2630 KOps/s 2.2938 KOps/s $\color{#d91a1a}-1.34\%$
test_compile_add_self_flat[pytree-compile] 0.4729ms 0.1922ms 5.2018 KOps/s 5.0092 KOps/s $\color{#35bf28}+3.85\%$
test_compile_copy_flat[tensordict-compile] 53.2710μs 19.0445μs 52.5087 KOps/s 49.9869 KOps/s $\textbf{\color{#35bf28}+5.04\%}$
test_compile_copy_flat[tensordict-eager] 69.0110μs 41.5197μs 24.0850 KOps/s 24.1119 KOps/s $\color{#d91a1a}-0.11\%$
test_compile_copy_flat[pytree-compile] 0.2149ms 21.3501μs 46.8381 KOps/s 45.4434 KOps/s $\color{#35bf28}+3.07\%$
test_compile_copy_flat[pytree-eager] 0.3342ms 52.1173μs 19.1875 KOps/s 18.9571 KOps/s $\color{#35bf28}+1.22\%$
test_compile_assign_and_add[tensordict-compile] 2.2161ms 0.2363ms 4.2315 KOps/s 3.9314 KOps/s $\textbf{\color{#35bf28}+7.63\%}$
test_compile_assign_and_add[tensordict-eager] 3.5372ms 3.3059ms 302.4866 Ops/s 305.1189 Ops/s $\color{#d91a1a}-0.86\%$
test_compile_assign_and_add[pytree-compile] 2.1185ms 0.1800ms 5.5568 KOps/s 5.2730 KOps/s $\textbf{\color{#35bf28}+5.38\%}$
test_compile_assign_and_add[pytree-eager] 3.0076ms 2.8554ms 350.2158 Ops/s 361.2389 Ops/s $\color{#d91a1a}-3.05\%$
test_compile_indexing[tensor-tensordict-compile] 0.2351ms 0.1333ms 7.5039 KOps/s 6.9457 KOps/s $\textbf{\color{#35bf28}+8.04\%}$
test_compile_indexing[tensor-tensordict-eager] 0.3145ms 77.9620μs 12.8268 KOps/s 12.6970 KOps/s $\color{#35bf28}+1.02\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2368ms 0.1183ms 8.4548 KOps/s 8.2429 KOps/s $\color{#35bf28}+2.57\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2579ms 47.2151μs 21.1797 KOps/s 22.1481 KOps/s $\color{#d91a1a}-4.37\%$
test_compile_indexing[tensor-pytree-compile] 0.2223ms 0.1189ms 8.4091 KOps/s 8.2829 KOps/s $\color{#35bf28}+1.52\%$
test_compile_indexing[tensor-pytree-eager] 0.4696ms 46.7036μs 21.4116 KOps/s 22.1415 KOps/s $\color{#d91a1a}-3.30\%$
test_compile_indexing[slice-tensordict-compile] 0.2187ms 83.2947μs 12.0056 KOps/s 11.4359 KOps/s $\color{#35bf28}+4.98\%$
test_compile_indexing[slice-tensordict-eager] 0.2040ms 29.8483μs 33.5028 KOps/s 34.5995 KOps/s $\color{#d91a1a}-3.17\%$
test_compile_indexing[slice-tensorclass-compile] 0.1584ms 70.9619μs 14.0921 KOps/s 13.8084 KOps/s $\color{#35bf28}+2.05\%$
test_compile_indexing[slice-tensorclass-eager] 0.2406ms 22.8862μs 43.6945 KOps/s 44.9998 KOps/s $\color{#d91a1a}-2.90\%$
test_compile_indexing[slice-pytree-compile] 0.1255ms 71.3144μs 14.0224 KOps/s 13.5464 KOps/s $\color{#35bf28}+3.51\%$
test_compile_indexing[slice-pytree-eager] 0.2354ms 22.8974μs 43.6731 KOps/s 45.2644 KOps/s $\color{#d91a1a}-3.52\%$
test_compile_indexing[int-tensordict-compile] 0.1627ms 83.8524μs 11.9257 KOps/s 11.3343 KOps/s $\textbf{\color{#35bf28}+5.22\%}$
test_compile_indexing[int-tensordict-eager] 0.4487ms 27.9377μs 35.7939 KOps/s 35.6483 KOps/s $\color{#35bf28}+0.41\%$
test_compile_indexing[int-tensorclass-compile] 0.5166ms 71.2997μs 14.0253 KOps/s 13.4881 KOps/s $\color{#35bf28}+3.98\%$
test_compile_indexing[int-tensorclass-eager] 0.4567ms 22.7971μs 43.8653 KOps/s 44.8650 KOps/s $\color{#d91a1a}-2.23\%$
test_compile_indexing[int-pytree-compile] 0.5114ms 71.0621μs 14.0722 KOps/s 13.5563 KOps/s $\color{#35bf28}+3.81\%$
test_compile_indexing[int-pytree-eager] 0.4403ms 22.7176μs 44.0188 KOps/s 45.1515 KOps/s $\color{#d91a1a}-2.51\%$
test_compile_replace[single-eager] 0.1064ms 48.7150μs 20.5276 KOps/s 21.1456 KOps/s $\color{#d91a1a}-2.92\%$
test_compile_replace[single-compile] 0.2498ms 0.1295ms 7.7217 KOps/s 7.4053 KOps/s $\color{#35bf28}+4.27\%$
test_compile_replace[multi-eager] 1.0254ms 0.5737ms 1.7430 KOps/s 1.7057 KOps/s $\color{#35bf28}+2.19\%$
test_compile_replace[multi-compile] 0.5779ms 0.1381ms 7.2418 KOps/s 7.1945 KOps/s $\color{#35bf28}+0.66\%$
test_compile_tc_getattr_20[eager] 0.6025ms 0.1750ms 5.7153 KOps/s 6.1883 KOps/s $\textbf{\color{#d91a1a}-7.64\%}$
test_compile_tc_getattr_20[compile] 0.4201ms 0.1462ms 6.8389 KOps/s 6.5963 KOps/s $\color{#35bf28}+3.68\%$
test_compile_clone_shallow[20-eager] 0.4541ms 19.3398μs 51.7067 KOps/s 51.9447 KOps/s $\color{#d91a1a}-0.46\%$
test_compile_clone_shallow[20-compile] 0.4569ms 16.6932μs 59.9047 KOps/s 59.7127 KOps/s $\color{#35bf28}+0.32\%$
test_compile_clone_shallow[40-eager] 0.4522ms 34.2876μs 29.1651 KOps/s 29.5167 KOps/s $\color{#d91a1a}-1.19\%$
test_compile_clone_shallow[40-compile] 0.4497ms 18.2035μs 54.9344 KOps/s 53.1566 KOps/s $\color{#35bf28}+3.34\%$
test_compile_clone_shallow[80-eager] 0.4950ms 63.4925μs 15.7499 KOps/s 16.1236 KOps/s $\color{#d91a1a}-2.32\%$
test_compile_clone_shallow[80-compile] 88.3510μs 21.3499μs 46.8386 KOps/s 47.4275 KOps/s $\color{#d91a1a}-1.24\%$
test_compile_update_inplace[eager] 0.4913ms 61.8769μs 16.1611 KOps/s 16.2748 KOps/s $\color{#d91a1a}-0.70\%$
test_compile_update_inplace[compile] 0.5858ms 0.1451ms 6.8931 KOps/s 6.6850 KOps/s $\color{#35bf28}+3.11\%$
test_mod_add[eager] 0.3010ms 0.2285ms 4.3755 KOps/s 4.4126 KOps/s $\color{#d91a1a}-0.84\%$
test_mod_add[compile] 0.5571ms 0.1315ms 7.6036 KOps/s 7.3365 KOps/s $\color{#35bf28}+3.64\%$
test_mod_add[compile-overhead] 0.2707ms 0.1774ms 5.6368 KOps/s 5.5162 KOps/s $\color{#35bf28}+2.19\%$
test_mod_wrap[eager] 0.9068ms 0.4726ms 2.1158 KOps/s 2.0872 KOps/s $\color{#35bf28}+1.37\%$
test_mod_wrap[compile] 0.5336ms 0.4142ms 2.4141 KOps/s 2.4060 KOps/s $\color{#35bf28}+0.34\%$
test_mod_wrap[compile-overhead] 8.3682ms 4.6371ms 215.6513 Ops/s 213.1267 Ops/s $\color{#35bf28}+1.18\%$
test_mod_wrap_and_backward[eager] 2.3247ms 1.7178ms 582.1544 Ops/s 544.7027 Ops/s $\textbf{\color{#35bf28}+6.88\%}$
test_mod_wrap_and_backward[compile] 2.4431ms 1.6192ms 617.5845 Ops/s 648.9233 Ops/s $\color{#d91a1a}-4.83\%$
test_mod_wrap_and_backward[compile-overhead] 1.5557ms 1.0537ms 949.0320 Ops/s 1.0324 KOps/s $\textbf{\color{#d91a1a}-8.07\%}$
test_seq_add[eager] 0.5769ms 0.5105ms 1.9587 KOps/s 1.9507 KOps/s $\color{#35bf28}+0.41\%$
test_seq_add[compile] 0.1880ms 0.1400ms 7.1416 KOps/s 6.8168 KOps/s $\color{#35bf28}+4.77\%$
test_seq_add[compile-overhead] 0.3072ms 0.1844ms 5.4243 KOps/s 5.3287 KOps/s $\color{#35bf28}+1.79\%$
test_seq_wrap[eager] 1.5144ms 1.4247ms 701.9166 Ops/s 708.4484 Ops/s $\color{#d91a1a}-0.92\%$
test_seq_wrap[compile] 0.4792ms 0.4164ms 2.4013 KOps/s 2.3761 KOps/s $\color{#35bf28}+1.06\%$
test_seq_wrap[compile-overhead] 0.3840ms 0.3058ms 3.2703 KOps/s 3.1884 KOps/s $\color{#35bf28}+2.57\%$
test_func_call_runtime[False-eager] 0.9373ms 0.8565ms 1.1676 KOps/s 1.2055 KOps/s $\color{#d91a1a}-3.15\%$
test_func_call_runtime[False-compile] 1.0500ms 0.9853ms 1.0150 KOps/s 1.0160 KOps/s $\color{#d91a1a}-0.10\%$
test_func_call_runtime[False-compile-overhead] 0.6750ms 0.5115ms 1.9550 KOps/s 1.9301 KOps/s $\color{#35bf28}+1.29\%$
test_func_call_runtime[True-eager] 1.1638ms 1.0507ms 951.7089 Ops/s 950.5190 Ops/s $\color{#35bf28}+0.13\%$
test_func_call_runtime[True-compile] 1.1237ms 1.0074ms 992.6913 Ops/s 995.8525 Ops/s $\color{#d91a1a}-0.32\%$
test_func_call_runtime[True-compile-overhead] 0.5745ms 0.5253ms 1.9036 KOps/s 1.8540 KOps/s $\color{#35bf28}+2.67\%$
test_func_call_cm_runtime[False-eager] 0.8964ms 0.8039ms 1.2439 KOps/s 1.2589 KOps/s $\color{#d91a1a}-1.19\%$
test_func_call_cm_runtime[False-compile] 1.1252ms 0.9722ms 1.0286 KOps/s 1.0328 KOps/s $\color{#d91a1a}-0.41\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5867ms 0.4928ms 2.0292 KOps/s 2.0025 KOps/s $\color{#35bf28}+1.33\%$
test_func_call_cm_runtime[True-eager] 1.3183ms 1.1998ms 833.4654 Ops/s 841.5695 Ops/s $\color{#d91a1a}-0.96\%$
test_func_call_cm_runtime[True-compile] 1.0840ms 1.0283ms 972.4570 Ops/s 976.8224 Ops/s $\color{#d91a1a}-0.45\%$
test_func_call_cm_runtime[True-compile-overhead] 0.7345ms 0.5487ms 1.8225 KOps/s 1.8113 KOps/s $\color{#35bf28}+0.62\%$
test_vmap_func_call_cm_runtime[eager] 2.8716ms 2.3847ms 419.3399 Ops/s 416.8757 Ops/s $\color{#35bf28}+0.59\%$
test_vmap_func_call_cm_runtime[compile] 1.1217ms 1.0350ms 966.1762 Ops/s 968.3479 Ops/s $\color{#d91a1a}-0.22\%$
test_vmap_func_call_cm_runtime[compile-overhead] 1.0756ms 0.5397ms 1.8529 KOps/s 1.8290 KOps/s $\color{#35bf28}+1.31\%$
test_distributed 1.6601ms 0.1553ms 6.4392 KOps/s 6.5492 KOps/s $\color{#d91a1a}-1.68\%$
test_tdmodule 0.4701ms 0.1934ms 5.1705 KOps/s 5.0599 KOps/s $\color{#35bf28}+2.19\%$
test_tdmodule_dispatch 0.4160ms 0.3710ms 2.6952 KOps/s 2.6696 KOps/s $\color{#35bf28}+0.96\%$
test_tdseq 39.8110μs 26.4346μs 37.8293 KOps/s 36.4341 KOps/s $\color{#35bf28}+3.83\%$
test_tdseq_dispatch 0.2410ms 0.2151ms 4.6496 KOps/s 4.6137 KOps/s $\color{#35bf28}+0.78\%$
test_instantiation_functorch 2.1558ms 2.0748ms 481.9772 Ops/s 481.9250 Ops/s $\color{#35bf28}+0.01\%$
test_exec_functorch 0.2316ms 0.1850ms 5.4057 KOps/s 5.5897 KOps/s $\color{#d91a1a}-3.29\%$
test_exec_functional_call 0.2378ms 0.1645ms 6.0784 KOps/s 6.1181 KOps/s $\color{#d91a1a}-0.65\%$
test_exec_td_decorator 0.4659ms 0.2413ms 4.1440 KOps/s 4.1402 KOps/s $\color{#35bf28}+0.09\%$
test_vmap_mlp_speed_decorator[True-True] 1.1921ms 1.0160ms 984.2067 Ops/s 978.1375 Ops/s $\color{#35bf28}+0.62\%$
test_vmap_mlp_speed_decorator[True-False] 1.1916ms 1.0133ms 986.9023 Ops/s 987.6041 Ops/s $\color{#d91a1a}-0.07\%$
test_vmap_mlp_speed_decorator[False-True] 0.8912ms 0.7137ms 1.4011 KOps/s 1.3913 KOps/s $\color{#35bf28}+0.70\%$
test_vmap_mlp_speed_decorator[False-False] 0.8929ms 0.7156ms 1.3974 KOps/s 1.3948 KOps/s $\color{#35bf28}+0.19\%$
test_vmap_transformer_speed_decorator[True-True] 22.1353ms 21.1331ms 47.3190 Ops/s 47.5849 Ops/s $\color{#d91a1a}-0.56\%$
test_vmap_transformer_speed_decorator[True-False] 21.8375ms 21.0266ms 47.5589 Ops/s 47.6034 Ops/s $\color{#d91a1a}-0.09\%$
test_vmap_transformer_speed_decorator[False-True] 21.7862ms 20.7842ms 48.1134 Ops/s 48.4332 Ops/s $\color{#d91a1a}-0.66\%$
test_vmap_transformer_speed_decorator[False-False] 21.3842ms 20.6394ms 48.4511 Ops/s 48.3525 Ops/s $\color{#35bf28}+0.20\%$
test_to_module_speed[True] 1.6216ms 1.4825ms 674.5369 Ops/s 672.3497 Ops/s $\color{#35bf28}+0.33\%$
test_to_module_speed[False] 1.5653ms 1.4448ms 692.1141 Ops/s 683.3982 Ops/s $\color{#35bf28}+1.28\%$
test_tc_init 68.3310μs 42.6631μs 23.4394 KOps/s 22.8750 KOps/s $\color{#35bf28}+2.47\%$
test_tc_init_tensor_only 42.6610μs 9.1634μs 109.1303 KOps/s 110.4878 KOps/s $\color{#d91a1a}-1.23\%$
test_tc_init_nested 0.1336ms 85.2612μs 11.7287 KOps/s 11.7639 KOps/s $\color{#d91a1a}-0.30\%$
test_tc_init_many_fields 35.3200μs 14.8297μs 67.4321 KOps/s 68.4574 KOps/s $\color{#d91a1a}-1.50\%$
test_tc_first_layer_tensor 18.4700μs 1.8237μs 548.3497 KOps/s 543.9636 KOps/s $\color{#35bf28}+0.81\%$
test_tc_first_layer_tensor_only 1.7300μs 0.3998μs 2.5010 MOps/s 2.4667 MOps/s $\color{#35bf28}+1.39\%$
test_tc_first_layer_tensor_set 24.9810μs 3.9230μs 254.9079 KOps/s 254.6064 KOps/s $\color{#35bf28}+0.12\%$
test_tc_first_layer_tensor_only_set 31.8810μs 3.2429μs 308.3681 KOps/s 306.9086 KOps/s $\color{#35bf28}+0.48\%$
test_tc_first_layer_nontensor 26.6610μs 6.1565μs 162.4299 KOps/s 155.7442 KOps/s $\color{#35bf28}+4.29\%$
test_tc_second_layer_tensor 50.7310μs 4.4132μs 226.5905 KOps/s 226.8246 KOps/s $\color{#d91a1a}-0.10\%$
test_tc_second_layer_nontensor 31.6100μs 8.7406μs 114.4086 KOps/s 111.8743 KOps/s $\color{#35bf28}+2.27\%$
test_unbind 0.2805s 13.8980ms 71.9525 Ops/s 68.7212 Ops/s $\color{#35bf28}+4.70\%$
test_full_like 4.5475ms 4.4027ms 227.1333 Ops/s 227.4981 Ops/s $\color{#d91a1a}-0.16\%$
test_zeros_like 4.5047ms 4.3762ms 228.5091 Ops/s 228.8464 Ops/s $\color{#d91a1a}-0.15\%$
test_ones_like 4.5445ms 4.3953ms 227.5167 Ops/s 228.4746 Ops/s $\color{#d91a1a}-0.42\%$
test_clone 6.8624ms 6.6033ms 151.4390 Ops/s 152.0166 Ops/s $\color{#d91a1a}-0.38\%$
test_squeeze 61.4410μs 14.2836μs 70.0102 KOps/s 71.1477 KOps/s $\color{#d91a1a}-1.60\%$
test_unsqueeze 0.1647ms 0.1147ms 8.7193 KOps/s 8.9402 KOps/s $\color{#d91a1a}-2.47\%$
test_split 0.2401ms 0.1795ms 5.5706 KOps/s 5.5529 KOps/s $\color{#35bf28}+0.32\%$
test_permute 0.2560ms 0.2119ms 4.7196 KOps/s 4.7547 KOps/s $\color{#d91a1a}-0.74\%$
test_stack 52.0718ms 51.8015ms 19.3045 Ops/s 36.5654 Ops/s $\textbf{\color{#d91a1a}-47.21\%}$
test_cat 52.1311ms 51.7190ms 19.3352 Ops/s 36.7403 Ops/s $\textbf{\color{#d91a1a}-47.37\%}$
test_sequential_tensordict 0.6597ms 0.5826ms 1.7163 KOps/s 1.7263 KOps/s $\color{#d91a1a}-0.57\%$
test_sequential_graph_module 0.1878ms 0.1253ms 7.9811 KOps/s 8.0692 KOps/s $\color{#d91a1a}-1.09\%$
test_nested_tensordict 0.9773ms 0.8233ms 1.2147 KOps/s 1.2229 KOps/s $\color{#d91a1a}-0.67\%$
test_nested_graph_module 0.1939ms 0.1365ms 7.3246 KOps/s 7.3849 KOps/s $\color{#d91a1a}-0.82\%$

@github-actions

github-actions Bot commented May 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 341. Improved: $\large\color{#35bf28}19$. Worsened: $\large\color{#d91a1a}22$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_collector_preallocated_write[stack_out-cpu-32-2] 0.2215ms 0.1718ms 5.8195 KOps/s 5.9077 KOps/s $\color{#d91a1a}-1.49\%$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.3816ms 0.2993ms 3.3413 KOps/s 3.3728 KOps/s $\color{#d91a1a}-0.93\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8444ms 0.7817ms 1.2793 KOps/s 1.2821 KOps/s $\color{#d91a1a}-0.22\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.4598ms 0.3420ms 2.9239 KOps/s 2.8346 KOps/s $\color{#35bf28}+3.15\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.4062ms 1.1449ms 873.4019 Ops/s 909.1365 Ops/s $\color{#d91a1a}-3.93\%$
test_collector_preallocated_write[stack_out-cpu-4096-32] 7.4937ms 7.3312ms 136.4028 Ops/s 134.9507 Ops/s $\color{#35bf28}+1.08\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.3014ms 0.2567ms 3.8949 KOps/s 3.9543 KOps/s $\color{#d91a1a}-1.50\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.4679ms 0.3768ms 2.6541 KOps/s 2.6540 KOps/s $+0.00\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.8972ms 0.8327ms 1.2010 KOps/s 1.2034 KOps/s $\color{#d91a1a}-0.20\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.3373ms 0.2541ms 3.9350 KOps/s 3.7938 KOps/s $\color{#35bf28}+3.72\%$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.4385ms 0.3783ms 2.6433 KOps/s 2.5768 KOps/s $\color{#35bf28}+2.58\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 0.9392ms 0.8853ms 1.1296 KOps/s 1.1294 KOps/s $\color{#35bf28}+0.02\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.3008ms 0.2290ms 4.3667 KOps/s 4.3997 KOps/s $\color{#d91a1a}-0.75\%$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 1.0059ms 0.8917ms 1.1215 KOps/s 1.0961 KOps/s $\color{#35bf28}+2.31\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.7880ms 3.5797ms 279.3540 Ops/s 302.6743 Ops/s $\textbf{\color{#d91a1a}-7.70\%}$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.9215ms 0.3767ms 2.6543 KOps/s 2.6439 KOps/s $\color{#35bf28}+0.40\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 1.9480ms 1.7676ms 565.7373 Ops/s 574.9555 Ops/s $\color{#d91a1a}-1.60\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 10.8347ms 10.3385ms 96.7254 Ops/s 92.9873 Ops/s $\color{#35bf28}+4.02\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.4696ms 0.4000ms 2.5000 KOps/s 2.5245 KOps/s $\color{#d91a1a}-0.97\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.6719ms 1.5717ms 636.2412 Ops/s 624.9222 Ops/s $\color{#35bf28}+1.81\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 6.2957ms 6.1157ms 163.5123 Ops/s 159.0028 Ops/s $\color{#35bf28}+2.84\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.4763ms 0.4177ms 2.3941 KOps/s 2.3424 KOps/s $\color{#35bf28}+2.21\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.6827ms 1.5610ms 640.6160 Ops/s 623.7698 Ops/s $\color{#35bf28}+2.70\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 6.3735ms 6.1085ms 163.7075 Ops/s 157.5755 Ops/s $\color{#35bf28}+3.89\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.2286ms 0.1725ms 5.7961 KOps/s 5.6568 KOps/s $\color{#35bf28}+2.46\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 0.7313ms 0.6709ms 1.4906 KOps/s 1.4435 KOps/s $\color{#35bf28}+3.27\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 2.8146ms 2.6691ms 374.6628 Ops/s 372.0709 Ops/s $\color{#35bf28}+0.70\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.4438ms 0.3124ms 3.2006 KOps/s 3.0654 KOps/s $\color{#35bf28}+4.41\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 2.0639ms 1.6055ms 622.8530 Ops/s 643.2714 Ops/s $\color{#d91a1a}-3.17\%$
test_collector_preallocated_write[copy_at-cpu-4096-32] 11.3171ms 9.5600ms 104.6025 Ops/s 104.0207 Ops/s $\color{#35bf28}+0.56\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.3819ms 0.3249ms 3.0780 KOps/s 3.0226 KOps/s $\color{#35bf28}+1.84\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.2874ms 1.2120ms 825.0667 Ops/s 821.5041 Ops/s $\color{#35bf28}+0.43\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 4.8604ms 4.7606ms 210.0574 Ops/s 210.9899 Ops/s $\color{#d91a1a}-0.44\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.4732ms 0.3266ms 3.0617 KOps/s 3.0518 KOps/s $\color{#35bf28}+0.32\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.3017ms 1.2160ms 822.3562 Ops/s 819.3550 Ops/s $\color{#35bf28}+0.37\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 4.8962ms 4.7819ms 209.1239 Ops/s 208.5359 Ops/s $\color{#35bf28}+0.28\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1544ms 0.1207ms 8.2850 KOps/s 8.2148 KOps/s $\color{#35bf28}+0.85\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5496ms 0.4743ms 2.1086 KOps/s 2.0933 KOps/s $\color{#35bf28}+0.73\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 2.0029ms 1.8940ms 527.9900 Ops/s 524.7723 Ops/s $\color{#35bf28}+0.61\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.6548ms 0.2608ms 3.8340 KOps/s 3.7259 KOps/s $\color{#35bf28}+2.90\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.5868ms 1.3496ms 740.9411 Ops/s 793.3131 Ops/s $\textbf{\color{#d91a1a}-6.60\%}$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 9.0589ms 8.6079ms 116.1718 Ops/s 115.9199 Ops/s $\color{#35bf28}+0.22\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.3287ms 0.2773ms 3.6068 KOps/s 3.3984 KOps/s $\textbf{\color{#35bf28}+6.13\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.1766ms 1.0501ms 952.2624 Ops/s 928.3654 Ops/s $\color{#35bf28}+2.57\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.4595ms 4.3315ms 230.8695 Ops/s 235.0249 Ops/s $\color{#d91a1a}-1.77\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.3783ms 0.2782ms 3.5950 KOps/s 3.3111 KOps/s $\textbf{\color{#35bf28}+8.57\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.1673ms 1.1146ms 897.1436 Ops/s 896.6413 Ops/s $\color{#35bf28}+0.06\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.4782ms 4.3666ms 229.0085 Ops/s 248.0862 Ops/s $\textbf{\color{#d91a1a}-7.69\%}$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 16.7285ms 16.5310ms 60.4923 Ops/s 61.2970 Ops/s $\color{#d91a1a}-1.31\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 33.1675ms 32.9320ms 30.3656 Ops/s 30.5275 Ops/s $\color{#d91a1a}-0.53\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 86.0766ms 85.3325ms 11.7189 Ops/s 11.8967 Ops/s $\color{#d91a1a}-1.49\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2025s 0.2018s 4.9553 Ops/s 5.1144 Ops/s $\color{#d91a1a}-3.11\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 13.9382ms 13.7803ms 72.5675 Ops/s 72.2286 Ops/s $\color{#35bf28}+0.47\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 27.1611ms 26.9762ms 37.0697 Ops/s 37.0060 Ops/s $\color{#35bf28}+0.17\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 14.1603ms 14.0243ms 71.3046 Ops/s 71.4840 Ops/s $\color{#d91a1a}-0.25\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 27.6041ms 27.3047ms 36.6237 Ops/s 36.3222 Ops/s $\color{#35bf28}+0.83\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 57.0839ms 56.6414ms 17.6549 Ops/s 17.9530 Ops/s $\color{#d91a1a}-1.66\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1142s 0.1128s 8.8677 Ops/s 9.4025 Ops/s $\textbf{\color{#d91a1a}-5.69\%}$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1502s 0.1490s 6.7101 Ops/s 6.7975 Ops/s $\color{#d91a1a}-1.29\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3227s 0.3214s 3.1115 Ops/s 3.0451 Ops/s $\color{#35bf28}+2.18\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 79.5782ms 78.0149ms 12.8181 Ops/s 12.3547 Ops/s $\color{#35bf28}+3.75\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1638s 0.1629s 6.1375 Ops/s 6.6669 Ops/s $\textbf{\color{#d91a1a}-7.94\%}$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 77.4299ms 76.6234ms 13.0508 Ops/s 13.2292 Ops/s $\color{#d91a1a}-1.35\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1648s 0.1624s 6.1588 Ops/s 6.3117 Ops/s $\color{#d91a1a}-2.42\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.5666ms 30.3409ms 32.9588 Ops/s 32.9021 Ops/s $\color{#35bf28}+0.17\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 61.7182ms 61.4753ms 16.2667 Ops/s 16.2282 Ops/s $\color{#35bf28}+0.24\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1213s 0.1200s 8.3314 Ops/s 8.5027 Ops/s $\color{#d91a1a}-2.01\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2685s 0.2675s 3.7383 Ops/s 3.6949 Ops/s $\color{#35bf28}+1.17\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 43.9633ms 43.6316ms 22.9192 Ops/s 23.0570 Ops/s $\color{#d91a1a}-0.60\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 87.6993ms 87.4592ms 11.4339 Ops/s 11.5652 Ops/s $\color{#d91a1a}-1.14\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 44.1229ms 43.8879ms 22.7853 Ops/s 22.9267 Ops/s $\color{#d91a1a}-0.62\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 88.0681ms 87.9352ms 11.3720 Ops/s 11.4555 Ops/s $\color{#d91a1a}-0.73\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 20.7352ms 20.5598ms 48.6387 Ops/s 48.7467 Ops/s $\color{#d91a1a}-0.22\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 42.1082ms 41.8908ms 23.8716 Ops/s 23.9184 Ops/s $\color{#d91a1a}-0.20\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1089s 0.1085s 9.2156 Ops/s 9.4155 Ops/s $\color{#d91a1a}-2.12\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2448s 0.2439s 4.0992 Ops/s 4.1492 Ops/s $\color{#d91a1a}-1.21\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 37.8686ms 37.3768ms 26.7546 Ops/s 27.4126 Ops/s $\color{#d91a1a}-2.40\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 74.8516ms 74.0065ms 13.5123 Ops/s 13.7792 Ops/s $\color{#d91a1a}-1.94\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 37.5493ms 36.9889ms 27.0351 Ops/s 27.3795 Ops/s $\color{#d91a1a}-1.26\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 74.6766ms 73.7213ms 13.5646 Ops/s 13.6666 Ops/s $\color{#d91a1a}-0.75\%$
test_plain_set_nested 46.5730μs 15.5659μs 64.2428 KOps/s 64.4362 KOps/s $\color{#d91a1a}-0.30\%$
test_plain_set_stack_nested 42.3920μs 15.3911μs 64.9726 KOps/s 64.6214 KOps/s $\color{#35bf28}+0.54\%$
test_plain_set_nested_inplace 45.3230μs 17.0916μs 58.5083 KOps/s 58.5235 KOps/s $\color{#d91a1a}-0.03\%$
test_plain_set_stack_nested_inplace 49.4630μs 16.7624μs 59.6575 KOps/s 59.0031 KOps/s $\color{#35bf28}+1.11\%$
test_items 24.2910μs 6.0368μs 165.6495 KOps/s 166.0500 KOps/s $\color{#d91a1a}-0.24\%$
test_items_nested 0.5555ms 0.4632ms 2.1591 KOps/s 2.1352 KOps/s $\color{#35bf28}+1.12\%$
test_items_nested_locked 0.5287ms 0.4706ms 2.1249 KOps/s 2.1208 KOps/s $\color{#35bf28}+0.20\%$
test_items_nested_leaf 0.1420ms 97.4365μs 10.2631 KOps/s 10.2462 KOps/s $\color{#35bf28}+0.17\%$
test_items_stack_nested 0.5474ms 0.4641ms 2.1546 KOps/s 2.1425 KOps/s $\color{#35bf28}+0.57\%$
test_items_stack_nested_leaf 0.1300ms 97.2330μs 10.2846 KOps/s 10.2348 KOps/s $\color{#35bf28}+0.49\%$
test_items_stack_nested_locked 0.5292ms 0.4693ms 2.1307 KOps/s 2.1304 KOps/s $\color{#35bf28}+0.01\%$
test_keys 26.5610μs 4.2485μs 235.3785 KOps/s 235.9170 KOps/s $\color{#d91a1a}-0.23\%$
test_keys_nested 0.1639ms 0.1306ms 7.6582 KOps/s 7.5815 KOps/s $\color{#35bf28}+1.01\%$
test_keys_nested_locked 0.7692ms 0.1390ms 7.1947 KOps/s 7.1067 KOps/s $\color{#35bf28}+1.24\%$
test_keys_nested_leaf 0.1509ms 0.1214ms 8.2395 KOps/s 8.0976 KOps/s $\color{#35bf28}+1.75\%$
test_keys_stack_nested 0.1819ms 0.1309ms 7.6380 KOps/s 7.4869 KOps/s $\color{#35bf28}+2.02\%$
test_keys_stack_nested_leaf 0.1676ms 0.1199ms 8.3427 KOps/s 8.1895 KOps/s $\color{#35bf28}+1.87\%$
test_keys_stack_nested_locked 0.1955ms 0.1396ms 7.1649 KOps/s 7.1457 KOps/s $\color{#35bf28}+0.27\%$
test_values 5.4382μs 1.0281μs 972.7016 KOps/s 972.5747 KOps/s $\color{#35bf28}+0.01\%$
test_values_nested 77.8540μs 52.7023μs 18.9745 KOps/s 18.7085 KOps/s $\color{#35bf28}+1.42\%$
test_values_nested_locked 83.9440μs 58.2116μs 17.1787 KOps/s 17.4798 KOps/s $\color{#d91a1a}-1.72\%$
test_values_nested_leaf 90.2440μs 60.2712μs 16.5917 KOps/s 16.1775 KOps/s $\color{#35bf28}+2.56\%$
test_values_stack_nested 83.1050μs 52.5803μs 19.0185 KOps/s 18.7363 KOps/s $\color{#35bf28}+1.51\%$
test_values_stack_nested_leaf 0.1253ms 59.8550μs 16.7070 KOps/s 16.3094 KOps/s $\color{#35bf28}+2.44\%$
test_values_stack_nested_locked 86.6140μs 58.1438μs 17.1987 KOps/s 17.5452 KOps/s $\color{#d91a1a}-1.97\%$
test_membership 5.2220μs 0.8492μs 1.1776 MOps/s 1.1785 MOps/s $\color{#d91a1a}-0.07\%$
test_membership_nested 30.4020μs 3.0531μs 327.5404 KOps/s 331.3831 KOps/s $\color{#d91a1a}-1.16\%$
test_membership_nested_leaf 33.2510μs 3.1018μs 322.3975 KOps/s 329.5745 KOps/s $\color{#d91a1a}-2.18\%$
test_membership_stacked_nested 38.5620μs 2.8809μs 347.1106 KOps/s 340.7167 KOps/s $\color{#35bf28}+1.88\%$
test_membership_stacked_nested_leaf 35.9710μs 2.8908μs 345.9228 KOps/s 344.0187 KOps/s $\color{#35bf28}+0.55\%$
test_membership_nested_last 33.5510μs 4.6177μs 216.5566 KOps/s 220.2711 KOps/s $\color{#d91a1a}-1.69\%$
test_membership_nested_leaf_last 34.0420μs 4.6194μs 216.4787 KOps/s 219.3679 KOps/s $\color{#d91a1a}-1.32\%$
test_membership_stacked_nested_last 28.2120μs 4.4024μs 227.1476 KOps/s 228.6407 KOps/s $\color{#d91a1a}-0.65\%$
test_membership_stacked_nested_leaf_last 29.4520μs 4.3600μs 229.3575 KOps/s 226.4349 KOps/s $\color{#35bf28}+1.29\%$
test_nested_getleaf 53.4130μs 22.0151μs 45.4233 KOps/s 44.8206 KOps/s $\color{#35bf28}+1.34\%$
test_nested_get 83.7950μs 20.9490μs 47.7350 KOps/s 47.3425 KOps/s $\color{#35bf28}+0.83\%$
test_stacked_getleaf 47.5830μs 21.8192μs 45.8313 KOps/s 45.5061 KOps/s $\color{#35bf28}+0.71\%$
test_stacked_get 43.2820μs 20.7314μs 48.2360 KOps/s 47.7191 KOps/s $\color{#35bf28}+1.08\%$
test_nested_getitemleaf 51.2930μs 22.6044μs 44.2391 KOps/s 43.9955 KOps/s $\color{#35bf28}+0.55\%$
test_nested_getitem 52.0230μs 21.0776μs 47.4437 KOps/s 46.7754 KOps/s $\color{#35bf28}+1.43\%$
test_stacked_getitemleaf 47.9530μs 21.9721μs 45.5122 KOps/s 44.7536 KOps/s $\color{#35bf28}+1.70\%$
test_stacked_getitem 46.8420μs 21.1136μs 47.3629 KOps/s 46.7529 KOps/s $\color{#35bf28}+1.30\%$
test_lock_nested 0.5836ms 0.4782ms 2.0913 KOps/s 2.0626 KOps/s $\color{#35bf28}+1.39\%$
test_lock_stack_nested 0.5696ms 0.4830ms 2.0702 KOps/s 2.0783 KOps/s $\color{#d91a1a}-0.39\%$
test_unlock_nested 0.4575ms 0.3914ms 2.5552 KOps/s 2.5788 KOps/s $\color{#d91a1a}-0.92\%$
test_unlock_stack_nested 0.4724ms 0.3913ms 2.5556 KOps/s 2.5652 KOps/s $\color{#d91a1a}-0.38\%$
test_flatten_speed 0.1812ms 0.1230ms 8.1309 KOps/s 8.1151 KOps/s $\color{#35bf28}+0.19\%$
test_unflatten_speed 0.6643ms 0.5714ms 1.7502 KOps/s 1.7426 KOps/s $\color{#35bf28}+0.44\%$
test_common_ops 0.1001s 1.7825ms 560.9992 Ops/s 434.1650 Ops/s $\textbf{\color{#35bf28}+29.21\%}$
test_creation 72.7640μs 3.2075μs 311.7730 KOps/s 321.1509 KOps/s $\color{#d91a1a}-2.92\%$
test_creation_empty 23.6210μs 7.1363μs 140.1284 KOps/s 144.4669 KOps/s $\color{#d91a1a}-3.00\%$
test_creation_nested_1 28.5810μs 11.6516μs 85.8248 KOps/s 86.5454 KOps/s $\color{#d91a1a}-0.83\%$
test_creation_nested_2 36.0620μs 13.7933μs 72.4990 KOps/s 72.8219 KOps/s $\color{#d91a1a}-0.44\%$
test_creation_many_keys[10] 50.6120μs 21.3312μs 46.8797 KOps/s 47.5609 KOps/s $\color{#d91a1a}-1.43\%$
test_creation_many_keys[50] 0.1216ms 91.5073μs 10.9281 KOps/s 11.1668 KOps/s $\color{#d91a1a}-2.14\%$
test_creation_many_keys[100] 0.2305ms 0.1796ms 5.5681 KOps/s 5.7096 KOps/s $\color{#d91a1a}-2.48\%$
test_creation_nested_many_keys[10] 73.5240μs 45.9822μs 21.7475 KOps/s 22.0865 KOps/s $\color{#d91a1a}-1.53\%$
test_creation_nested_many_keys[50] 0.2635ms 0.1876ms 5.3292 KOps/s 5.4298 KOps/s $\color{#d91a1a}-1.85\%$
test_clone 42.8220μs 13.5679μs 73.7036 KOps/s 74.7947 KOps/s $\color{#d91a1a}-1.46\%$
test_getitem[int] 1.4816ms 15.2433μs 65.6027 KOps/s 65.6407 KOps/s $\color{#d91a1a}-0.06\%$
test_getitem[slice_int] 0.1591ms 24.4640μs 40.8764 KOps/s 40.2385 KOps/s $\color{#35bf28}+1.59\%$
test_getitem[range] 0.1990ms 65.4139μs 15.2873 KOps/s 14.5055 KOps/s $\textbf{\color{#35bf28}+5.39\%}$
test_getitem[tuple] 0.1399ms 24.3681μs 41.0373 KOps/s 40.8808 KOps/s $\color{#35bf28}+0.38\%$
test_getitem[list] 0.1890ms 61.0249μs 16.3868 KOps/s 16.1307 KOps/s $\color{#35bf28}+1.59\%$
test_setitem_dim[int] 49.7830μs 28.4429μs 35.1582 KOps/s 36.0386 KOps/s $\color{#d91a1a}-2.44\%$
test_setitem_dim[slice_int] 68.2040μs 45.2111μs 22.1185 KOps/s 21.7084 KOps/s $\color{#35bf28}+1.89\%$
test_setitem_dim[range] 0.1255ms 98.7633μs 10.1252 KOps/s 9.9005 KOps/s $\color{#35bf28}+2.27\%$
test_setitem_dim[tuple] 66.5840μs 43.2033μs 23.1464 KOps/s 24.4687 KOps/s $\textbf{\color{#d91a1a}-5.40\%}$
test_setitem 38.9620μs 18.1565μs 55.0767 KOps/s 56.3919 KOps/s $\color{#d91a1a}-2.33\%$
test_set 51.3930μs 17.3186μs 57.7413 KOps/s 58.8402 KOps/s $\color{#d91a1a}-1.87\%$
test_set_shared 0.8581ms 0.2131ms 4.6936 KOps/s 4.7152 KOps/s $\color{#d91a1a}-0.46\%$
test_update 0.5040ms 22.2434μs 44.9571 KOps/s 44.1404 KOps/s $\color{#35bf28}+1.85\%$
test_update_nested 71.6030μs 34.0750μs 29.3470 KOps/s 29.1275 KOps/s $\color{#35bf28}+0.75\%$
test_update__nested 0.4364ms 34.0734μs 29.3484 KOps/s 28.4025 KOps/s $\color{#35bf28}+3.33\%$
test_set_nested 70.5340μs 19.2106μs 52.0545 KOps/s 50.9089 KOps/s $\color{#35bf28}+2.25\%$
test_set_nested_new 54.7830μs 24.1401μs 41.4249 KOps/s 41.0112 KOps/s $\color{#35bf28}+1.01\%$
test_select 0.1174ms 40.7390μs 24.5465 KOps/s 23.7091 KOps/s $\color{#35bf28}+3.53\%$
test_select_nested 0.1047ms 74.4157μs 13.4380 KOps/s 13.4732 KOps/s $\color{#d91a1a}-0.26\%$
test_exclude_nested 0.1293ms 91.5115μs 10.9276 KOps/s 10.9140 KOps/s $\color{#35bf28}+0.12\%$
test_empty[True] 0.4614ms 0.4013ms 2.4917 KOps/s 2.4750 KOps/s $\color{#35bf28}+0.67\%$
test_empty[False] 7.9355μs 1.3656μs 732.2761 KOps/s 725.2550 KOps/s $\color{#35bf28}+0.97\%$
test_to 0.1153ms 72.7206μs 13.7513 KOps/s 14.3519 KOps/s $\color{#d91a1a}-4.19\%$
test_to_nonblocking 0.1200ms 71.4483μs 13.9961 KOps/s 13.8718 KOps/s $\color{#35bf28}+0.90\%$
test_unbind_speed 0.4090ms 0.3422ms 2.9221 KOps/s 2.9460 KOps/s $\color{#d91a1a}-0.81\%$
test_unbind_speed_stack0 0.4132ms 0.3379ms 2.9597 KOps/s 3.0044 KOps/s $\color{#d91a1a}-1.49\%$
test_unbind_speed_stack1 0.1141s 0.8540ms 1.1710 KOps/s 1.1485 KOps/s $\color{#35bf28}+1.95\%$
test_split 1.2510ms 1.1515ms 868.4039 Ops/s 706.2539 Ops/s $\textbf{\color{#35bf28}+22.96\%}$
test_chunk 0.1135s 1.2276ms 814.5821 Ops/s 913.3803 Ops/s $\textbf{\color{#d91a1a}-10.82\%}$
test_to_cpu_blocking 20.1086ms 19.7843ms 50.5450 Ops/s 50.4957 Ops/s $\color{#35bf28}+0.10\%$
test_to_cpu_global_sync 12.0561ms 11.8574ms 84.3358 Ops/s 83.9956 Ops/s $\color{#35bf28}+0.41\%$
test_to_cpu_event_sync 13.1910ms 12.8182ms 78.0140 Ops/s 77.8182 Ops/s $\color{#35bf28}+0.25\%$
test_to_cpu_default 13.0487ms 12.7826ms 78.2312 Ops/s 78.2084 Ops/s $\color{#35bf28}+0.03\%$
test_consolidate[False-None] 0.1431s 4.7532ms 210.3839 Ops/s 238.7247 Ops/s $\textbf{\color{#d91a1a}-11.87\%}$
test_consolidate[default-None] 3.6005ms 2.1542ms 464.2023 Ops/s 482.1262 Ops/s $\color{#d91a1a}-3.72\%$
test_consolidate[reduce-overhead-None] 0.1929s 2.4613ms 406.2856 Ops/s 503.9968 Ops/s $\textbf{\color{#d91a1a}-19.39\%}$
test_consolidate_njt[False-None] 8.9224ms 8.5671ms 116.7262 Ops/s 115.5424 Ops/s $\color{#35bf28}+1.02\%$
test_to[False-False-None] 2.4094ms 2.1888ms 456.8679 Ops/s 451.6723 Ops/s $\color{#35bf28}+1.15\%$
test_to[True-False-None] 2.1842ms 1.8728ms 533.9496 Ops/s 523.9609 Ops/s $\color{#35bf28}+1.91\%$
test_to[within-False-None] 6.2626ms 6.1139ms 163.5617 Ops/s 160.2823 Ops/s $\color{#35bf28}+2.05\%$
test_to[True-default-None] 11.0722ms 10.8499ms 92.1671 Ops/s 92.0788 Ops/s $\color{#35bf28}+0.10\%$
test_to_njt[False-False-None] 8.8234ms 8.6654ms 115.4017 Ops/s 115.8742 Ops/s $\color{#d91a1a}-0.41\%$
test_to_njt[True-False-None] 7.1968ms 7.0111ms 142.6300 Ops/s 143.3627 Ops/s $\color{#d91a1a}-0.51\%$
test_to_njt[within-False-None] 15.9973ms 15.8669ms 63.0242 Ops/s 64.1284 Ops/s $\color{#d91a1a}-1.72\%$
test_creation[device0] 0.3496ms 0.1153ms 8.6760 KOps/s 8.4712 KOps/s $\color{#35bf28}+2.42\%$
test_creation_from_tensor 0.3608ms 0.1130ms 8.8498 KOps/s 8.6685 KOps/s $\color{#35bf28}+2.09\%$
test_add_one[memmap_tensor0] 0.1603ms 6.7153μs 148.9134 KOps/s 147.0625 KOps/s $\color{#35bf28}+1.26\%$
test_contiguous[memmap_tensor0] 20.9910μs 0.6714μs 1.4893 MOps/s 2.2094 MOps/s $\textbf{\color{#d91a1a}-32.59\%}$
test_stack[memmap_tensor0] 93.3650μs 4.2916μs 233.0126 KOps/s 236.5194 KOps/s $\color{#d91a1a}-1.48\%$
test_memmaptd_index 1.2010ms 0.2758ms 3.6262 KOps/s 3.6925 KOps/s $\color{#d91a1a}-1.80\%$
test_memmaptd_index_astensor 0.5370ms 0.3783ms 2.6435 KOps/s 2.6664 KOps/s $\color{#d91a1a}-0.86\%$
test_memmaptd_index_op 0.7850ms 0.6461ms 1.5478 KOps/s 1.5585 KOps/s $\color{#d91a1a}-0.69\%$
test_serialize_model 0.1382s 0.1343s 7.4457 Ops/s 7.3781 Ops/s $\color{#35bf28}+0.92\%$
test_serialize_model_pickle 1.3493s 1.2101s 0.8264 Ops/s 0.8268 Ops/s $\color{#d91a1a}-0.05\%$
test_serialize_weights 0.1358s 0.1341s 7.4594 Ops/s 7.4607 Ops/s $\color{#d91a1a}-0.02\%$
test_serialize_weights_returnearly 0.4590s 89.6676ms 11.1523 Ops/s 11.4826 Ops/s $\color{#d91a1a}-2.88\%$
test_serialize_weights_pickle 1.3672s 1.2158s 0.8225 Ops/s 0.8190 Ops/s $\color{#35bf28}+0.42\%$
test_reshape_pytree 0.2106ms 32.7528μs 30.5318 KOps/s 30.1451 KOps/s $\color{#35bf28}+1.28\%$
test_reshape_td 76.8740μs 44.6159μs 22.4135 KOps/s 22.2318 KOps/s $\color{#35bf28}+0.82\%$
test_view_pytree 0.2198ms 32.3840μs 30.8794 KOps/s 30.9266 KOps/s $\color{#d91a1a}-0.15\%$
test_view_td 94.0950μs 52.7317μs 18.9639 KOps/s 18.8314 KOps/s $\color{#35bf28}+0.70\%$
test_unbind_pytree 0.2340ms 36.7613μs 27.2025 KOps/s 27.3412 KOps/s $\color{#d91a1a}-0.51\%$
test_unbind_td 0.1125ms 50.5535μs 19.7810 KOps/s 20.0676 KOps/s $\color{#d91a1a}-1.43\%$
test_split_pytree 0.2473ms 42.8930μs 23.3138 KOps/s 23.5416 KOps/s $\color{#d91a1a}-0.97\%$
test_split_td 0.2137ms 65.3919μs 15.2924 KOps/s 15.6939 KOps/s $\color{#d91a1a}-2.56\%$
test_add_pytree 0.2300ms 43.2083μs 23.1437 KOps/s 23.3512 KOps/s $\color{#d91a1a}-0.89\%$
test_add_td 94.6160μs 56.0872μs 17.8294 KOps/s 18.1425 KOps/s $\color{#d91a1a}-1.73\%$
test_compile_add_one_nested[tensordict-compile] 0.3511ms 0.2264ms 4.4171 KOps/s 4.2876 KOps/s $\color{#35bf28}+3.02\%$
test_compile_add_one_nested[tensordict-eager] 0.3426ms 0.2041ms 4.9007 KOps/s 4.9763 KOps/s $\color{#d91a1a}-1.52\%$
test_compile_add_one_nested[pytree-compile] 0.2729ms 0.1375ms 7.2717 KOps/s 7.0144 KOps/s $\color{#35bf28}+3.67\%$
test_compile_add_one_nested[pytree-eager] 0.4370ms 0.1889ms 5.2938 KOps/s 5.4640 KOps/s $\color{#d91a1a}-3.11\%$
test_compile_copy_nested[tensordict-compile] 63.6930μs 15.5534μs 64.2946 KOps/s 67.8539 KOps/s $\textbf{\color{#d91a1a}-5.25\%}$
test_compile_copy_nested[tensordict-eager] 0.3202ms 54.0161μs 18.5130 KOps/s 18.2616 KOps/s $\color{#35bf28}+1.38\%$
test_compile_copy_nested[pytree-compile] 0.1356ms 16.1571μs 61.8924 KOps/s 63.1660 KOps/s $\color{#d91a1a}-2.02\%$
test_compile_copy_nested[pytree-eager] 0.3843ms 68.3057μs 14.6401 KOps/s 14.5498 KOps/s $\color{#35bf28}+0.62\%$
test_compile_add_one_flat[tensordict-compile] 0.3378ms 0.2549ms 3.9231 KOps/s 3.7122 KOps/s $\textbf{\color{#35bf28}+5.68\%}$
test_compile_add_one_flat[tensordict-eager] 0.4591ms 0.2818ms 3.5481 KOps/s 3.3949 KOps/s $\color{#35bf28}+4.51\%$
test_compile_add_one_flat[tensorclass-compile] 0.6507ms 0.1906ms 5.2457 KOps/s 4.8552 KOps/s $\textbf{\color{#35bf28}+8.04\%}$
test_compile_add_one_flat[tensorclass-eager] 0.5180ms 78.3330μs 12.7660 KOps/s 12.7556 KOps/s $\color{#35bf28}+0.08\%$
test_compile_add_one_flat[pytree-compile] 0.3019ms 0.2021ms 4.9473 KOps/s 4.8104 KOps/s $\color{#35bf28}+2.85\%$
test_compile_add_one_flat[pytree-eager] 0.9669ms 0.5292ms 1.8898 KOps/s 1.8525 KOps/s $\color{#35bf28}+2.01\%$
test_compile_add_self_flat[tensordict-eager] 0.7665ms 0.3303ms 3.0277 KOps/s 2.9497 KOps/s $\color{#35bf28}+2.64\%$
test_compile_add_self_flat[tensordict-compile] 0.4172ms 0.2547ms 3.9262 KOps/s 3.6705 KOps/s $\textbf{\color{#35bf28}+6.96\%}$
test_compile_add_self_flat[tensorclass-eager] 0.5323ms 92.7027μs 10.7872 KOps/s 10.5264 KOps/s $\color{#35bf28}+2.48\%$
test_compile_add_self_flat[tensorclass-compile] 0.3259ms 0.1885ms 5.3054 KOps/s 4.9069 KOps/s $\textbf{\color{#35bf28}+8.12\%}$
test_compile_add_self_flat[pytree-eager] 0.6542ms 0.4337ms 2.3056 KOps/s 2.2682 KOps/s $\color{#35bf28}+1.65\%$
test_compile_add_self_flat[pytree-compile] 0.2507ms 0.1965ms 5.0880 KOps/s 4.8564 KOps/s $\color{#35bf28}+4.77\%$
test_compile_copy_flat[tensordict-compile] 75.1140μs 19.8749μs 50.3147 KOps/s 50.8369 KOps/s $\color{#d91a1a}-1.03\%$
test_compile_copy_flat[tensordict-eager] 95.0250μs 42.2663μs 23.6595 KOps/s 24.0722 KOps/s $\color{#d91a1a}-1.71\%$
test_compile_copy_flat[pytree-compile] 0.2031ms 21.7617μs 45.9523 KOps/s 45.2369 KOps/s $\color{#35bf28}+1.58\%$
test_compile_copy_flat[pytree-eager] 0.4665ms 51.2370μs 19.5172 KOps/s 19.1221 KOps/s $\color{#35bf28}+2.07\%$
test_compile_assign_and_add[tensordict-compile] 2.2465ms 0.2398ms 4.1709 KOps/s 4.0084 KOps/s $\color{#35bf28}+4.06\%$
test_compile_assign_and_add[tensordict-eager] 3.4753ms 3.2828ms 304.6209 Ops/s 302.8075 Ops/s $\color{#35bf28}+0.60\%$
test_compile_assign_and_add[pytree-compile] 2.1350ms 0.1833ms 5.4543 KOps/s 5.4047 KOps/s $\color{#35bf28}+0.92\%$
test_compile_assign_and_add[pytree-eager] 2.9870ms 2.8003ms 357.1026 Ops/s 356.0933 Ops/s $\color{#35bf28}+0.28\%$
test_compile_indexing[tensor-tensordict-compile] 0.3323ms 0.1406ms 7.1116 KOps/s 7.2767 KOps/s $\color{#d91a1a}-2.27\%$
test_compile_indexing[tensor-tensordict-eager] 0.3233ms 76.9309μs 12.9987 KOps/s 13.1487 KOps/s $\color{#d91a1a}-1.14\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2234ms 0.1240ms 8.0658 KOps/s 8.0831 KOps/s $\color{#d91a1a}-0.21\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2629ms 45.6968μs 21.8833 KOps/s 21.9981 KOps/s $\color{#d91a1a}-0.52\%$
test_compile_indexing[tensor-pytree-compile] 0.5366ms 0.1211ms 8.2583 KOps/s 8.3488 KOps/s $\color{#d91a1a}-1.08\%$
test_compile_indexing[tensor-pytree-eager] 0.2875ms 45.4704μs 21.9923 KOps/s 20.9061 KOps/s $\textbf{\color{#35bf28}+5.20\%}$
test_compile_indexing[slice-tensordict-compile] 0.5146ms 85.6564μs 11.6746 KOps/s 11.9735 KOps/s $\color{#d91a1a}-2.50\%$
test_compile_indexing[slice-tensordict-eager] 0.4551ms 31.2024μs 32.0488 KOps/s 36.4994 KOps/s $\textbf{\color{#d91a1a}-12.19\%}$
test_compile_indexing[slice-tensorclass-compile] 0.1911ms 73.0726μs 13.6850 KOps/s 14.0392 KOps/s $\color{#d91a1a}-2.52\%$
test_compile_indexing[slice-tensorclass-eager] 0.4618ms 23.0558μs 43.3731 KOps/s 43.3045 KOps/s $\color{#35bf28}+0.16\%$
test_compile_indexing[slice-pytree-compile] 0.4996ms 75.0101μs 13.3315 KOps/s 14.0824 KOps/s $\textbf{\color{#d91a1a}-5.33\%}$
test_compile_indexing[slice-pytree-eager] 0.4567ms 23.2158μs 43.0740 KOps/s 44.7127 KOps/s $\color{#d91a1a}-3.66\%$
test_compile_indexing[int-tensordict-compile] 0.5267ms 89.1627μs 11.2155 KOps/s 11.8554 KOps/s $\textbf{\color{#d91a1a}-5.40\%}$
test_compile_indexing[int-tensordict-eager] 0.4775ms 29.5628μs 33.8263 KOps/s 36.5557 KOps/s $\textbf{\color{#d91a1a}-7.47\%}$
test_compile_indexing[int-tensorclass-compile] 0.5124ms 74.9717μs 13.3384 KOps/s 14.1558 KOps/s $\textbf{\color{#d91a1a}-5.77\%}$
test_compile_indexing[int-tensorclass-eager] 0.4751ms 22.9827μs 43.5109 KOps/s 44.2012 KOps/s $\color{#d91a1a}-1.56\%$
test_compile_indexing[int-pytree-compile] 0.1180ms 74.6030μs 13.4043 KOps/s 14.1329 KOps/s $\textbf{\color{#d91a1a}-5.16\%}$
test_compile_indexing[int-pytree-eager] 0.4496ms 22.8944μs 43.6788 KOps/s 43.8701 KOps/s $\color{#d91a1a}-0.44\%$
test_compile_replace[single-eager] 0.5072ms 50.0906μs 19.9638 KOps/s 19.8989 KOps/s $\color{#35bf28}+0.33\%$
test_compile_replace[single-compile] 0.5754ms 0.1364ms 7.3293 KOps/s 7.6758 KOps/s $\color{#d91a1a}-4.51\%$
test_compile_replace[multi-eager] 1.0934ms 0.5926ms 1.6874 KOps/s 1.7743 KOps/s $\color{#d91a1a}-4.90\%$
test_compile_replace[multi-compile] 0.5806ms 0.1446ms 6.9179 KOps/s 7.2346 KOps/s $\color{#d91a1a}-4.38\%$
test_compile_tc_getattr_20[eager] 0.6162ms 0.1660ms 6.0223 KOps/s 5.9150 KOps/s $\color{#35bf28}+1.81\%$
test_compile_tc_getattr_20[compile] 0.6033ms 0.1545ms 6.4708 KOps/s 6.7437 KOps/s $\color{#d91a1a}-4.05\%$
test_compile_clone_shallow[20-eager] 0.4576ms 19.3071μs 51.7943 KOps/s 51.1155 KOps/s $\color{#35bf28}+1.33\%$
test_compile_clone_shallow[20-compile] 80.5440μs 17.2588μs 57.9413 KOps/s 59.1401 KOps/s $\color{#d91a1a}-2.03\%$
test_compile_clone_shallow[40-eager] 0.4656ms 34.1664μs 29.2686 KOps/s 29.0088 KOps/s $\color{#35bf28}+0.90\%$
test_compile_clone_shallow[40-compile] 0.4705ms 19.3490μs 51.6822 KOps/s 53.6874 KOps/s $\color{#d91a1a}-3.73\%$
test_compile_clone_shallow[80-eager] 0.4957ms 63.4627μs 15.7573 KOps/s 15.4614 KOps/s $\color{#35bf28}+1.91\%$
test_compile_clone_shallow[80-compile] 76.7240μs 23.2959μs 42.9259 KOps/s 46.0693 KOps/s $\textbf{\color{#d91a1a}-6.82\%}$
test_compile_update_inplace[eager] 0.4905ms 60.9561μs 16.4052 KOps/s 16.3127 KOps/s $\color{#35bf28}+0.57\%$
test_compile_update_inplace[compile] 0.4187ms 0.1531ms 6.5322 KOps/s 6.5149 KOps/s $\color{#35bf28}+0.26\%$
test_mod_add[eager] 0.6705ms 0.2286ms 4.3746 KOps/s 4.3876 KOps/s $\color{#d91a1a}-0.30\%$
test_mod_add[compile] 0.2830ms 0.1318ms 7.5858 KOps/s 7.5012 KOps/s $\color{#35bf28}+1.13\%$
test_mod_add[compile-overhead] 0.3046ms 0.1807ms 5.5351 KOps/s 5.5589 KOps/s $\color{#d91a1a}-0.43\%$
test_mod_wrap[eager] 0.5662ms 0.4726ms 2.1161 KOps/s 2.1017 KOps/s $\color{#35bf28}+0.69\%$
test_mod_wrap[compile] 0.4696ms 0.3986ms 2.5088 KOps/s 2.4873 KOps/s $\color{#35bf28}+0.86\%$
test_mod_wrap[compile-overhead] 9.4098ms 5.1105ms 195.6761 Ops/s 214.7109 Ops/s $\textbf{\color{#d91a1a}-8.87\%}$
test_mod_wrap_and_backward[eager] 1.9392ms 1.8193ms 549.6624 Ops/s 585.7025 Ops/s $\textbf{\color{#d91a1a}-6.15\%}$
test_mod_wrap_and_backward[compile] 2.4454ms 1.5991ms 625.3700 Ops/s 635.5834 Ops/s $\color{#d91a1a}-1.61\%$
test_mod_wrap_and_backward[compile-overhead] 1.3212ms 0.9462ms 1.0569 KOps/s 1.0413 KOps/s $\color{#35bf28}+1.50\%$
test_seq_add[eager] 0.6079ms 0.5107ms 1.9582 KOps/s 1.9335 KOps/s $\color{#35bf28}+1.28\%$
test_seq_add[compile] 0.2529ms 0.1421ms 7.0356 KOps/s 6.6020 KOps/s $\textbf{\color{#35bf28}+6.57\%}$
test_seq_add[compile-overhead] 0.3067ms 0.1837ms 5.4436 KOps/s 5.3265 KOps/s $\color{#35bf28}+2.20\%$
test_seq_wrap[eager] 1.5201ms 1.4100ms 709.2310 Ops/s 704.0573 Ops/s $\color{#35bf28}+0.73\%$
test_seq_wrap[compile] 0.4882ms 0.4118ms 2.4281 KOps/s 2.3730 KOps/s $\color{#35bf28}+2.32\%$
test_seq_wrap[compile-overhead] 0.3775ms 0.3082ms 3.2446 KOps/s 3.0968 KOps/s $\color{#35bf28}+4.77\%$
test_func_call_runtime[False-eager] 0.8829ms 0.7987ms 1.2520 KOps/s 1.1653 KOps/s $\textbf{\color{#35bf28}+7.43\%}$
test_func_call_runtime[False-compile] 1.0620ms 0.9807ms 1.0197 KOps/s 961.4395 Ops/s $\textbf{\color{#35bf28}+6.06\%}$
test_func_call_runtime[False-compile-overhead] 0.6961ms 0.5195ms 1.9250 KOps/s 1.9308 KOps/s $\color{#d91a1a}-0.30\%$
test_func_call_runtime[True-eager] 1.1235ms 1.0487ms 953.5933 Ops/s 925.1472 Ops/s $\color{#35bf28}+3.07\%$
test_func_call_runtime[True-compile] 1.1040ms 0.9987ms 1.0013 KOps/s 993.6776 Ops/s $\color{#35bf28}+0.77\%$
test_func_call_runtime[True-compile-overhead] 0.6102ms 0.5344ms 1.8712 KOps/s 1.8796 KOps/s $\color{#d91a1a}-0.44\%$
test_func_call_cm_runtime[False-eager] 0.9181ms 0.8014ms 1.2479 KOps/s 1.2493 KOps/s $\color{#d91a1a}-0.12\%$
test_func_call_cm_runtime[False-compile] 1.0505ms 0.9618ms 1.0397 KOps/s 1.0252 KOps/s $\color{#35bf28}+1.42\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5990ms 0.4946ms 2.0218 KOps/s 2.0170 KOps/s $\color{#35bf28}+0.24\%$
test_func_call_cm_runtime[True-eager] 1.3024ms 1.1927ms 838.4516 Ops/s 838.3453 Ops/s $\color{#35bf28}+0.01\%$
test_func_call_cm_runtime[True-compile] 1.0888ms 1.0147ms 985.4678 Ops/s 978.8203 Ops/s $\color{#35bf28}+0.68\%$
test_func_call_cm_runtime[True-compile-overhead] 0.7085ms 0.5495ms 1.8200 KOps/s 1.8248 KOps/s $\color{#d91a1a}-0.26\%$
test_vmap_func_call_cm_runtime[eager] 2.8647ms 2.3848ms 419.3203 Ops/s 417.6869 Ops/s $\color{#35bf28}+0.39\%$
test_vmap_func_call_cm_runtime[compile] 1.1311ms 1.0249ms 975.6848 Ops/s 973.8935 Ops/s $\color{#35bf28}+0.18\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.5904ms 0.5367ms 1.8631 KOps/s 1.8426 KOps/s $\color{#35bf28}+1.12\%$
test_distributed 0.6857ms 0.1482ms 6.7462 KOps/s 6.6299 KOps/s $\color{#35bf28}+1.75\%$
test_tdmodule 0.6702ms 0.1936ms 5.1644 KOps/s 4.9879 KOps/s $\color{#35bf28}+3.54\%$
test_tdmodule_dispatch 0.4271ms 0.3707ms 2.6973 KOps/s 2.6170 KOps/s $\color{#35bf28}+3.07\%$
test_tdseq 45.4630μs 26.4422μs 37.8184 KOps/s 35.8691 KOps/s $\textbf{\color{#35bf28}+5.43\%}$
test_tdseq_dispatch 0.2718ms 0.2153ms 4.6443 KOps/s 4.4779 KOps/s $\color{#35bf28}+3.72\%$
test_instantiation_functorch 2.2627ms 2.0931ms 477.7609 Ops/s 480.5052 Ops/s $\color{#d91a1a}-0.57\%$
test_exec_functorch 0.2311ms 0.1823ms 5.4852 KOps/s 5.1720 KOps/s $\textbf{\color{#35bf28}+6.06\%}$
test_exec_functional_call 0.2141ms 0.1613ms 6.1984 KOps/s 5.7137 KOps/s $\textbf{\color{#35bf28}+8.48\%}$
test_exec_td_decorator 0.4531ms 0.2397ms 4.1717 KOps/s 4.0038 KOps/s $\color{#35bf28}+4.19\%$
test_vmap_mlp_speed_decorator[True-True] 1.2000ms 1.0116ms 988.5344 Ops/s 963.5070 Ops/s $\color{#35bf28}+2.60\%$
test_vmap_mlp_speed_decorator[True-False] 1.1890ms 1.0077ms 992.3522 Ops/s 951.3986 Ops/s $\color{#35bf28}+4.30\%$
test_vmap_mlp_speed_decorator[False-True] 0.8737ms 0.7147ms 1.3992 KOps/s 1.3603 KOps/s $\color{#35bf28}+2.86\%$
test_vmap_mlp_speed_decorator[False-False] 0.8926ms 0.7143ms 1.4000 KOps/s 1.3463 KOps/s $\color{#35bf28}+3.99\%$
test_vmap_transformer_speed_decorator[True-True] 21.7612ms 20.9582ms 47.7141 Ops/s 46.5559 Ops/s $\color{#35bf28}+2.49\%$
test_vmap_transformer_speed_decorator[True-False] 21.8050ms 21.0330ms 47.5442 Ops/s 47.1086 Ops/s $\color{#35bf28}+0.92\%$
test_vmap_transformer_speed_decorator[False-True] 21.2756ms 20.5715ms 48.6110 Ops/s 47.8065 Ops/s $\color{#35bf28}+1.68\%$
test_vmap_transformer_speed_decorator[False-False] 20.7064ms 20.4989ms 48.7830 Ops/s 48.1215 Ops/s $\color{#35bf28}+1.37\%$
test_to_module_speed[True] 1.5950ms 1.4867ms 672.6128 Ops/s 667.1440 Ops/s $\color{#35bf28}+0.82\%$
test_to_module_speed[False] 1.5616ms 1.4667ms 681.7819 Ops/s 684.4327 Ops/s $\color{#d91a1a}-0.39\%$
test_tc_init 69.8240μs 43.4621μs 23.0085 KOps/s 22.2364 KOps/s $\color{#35bf28}+3.47\%$
test_tc_init_tensor_only 26.5210μs 9.0999μs 109.8916 KOps/s 111.5866 KOps/s $\color{#d91a1a}-1.52\%$
test_tc_init_nested 0.1194ms 85.4053μs 11.7089 KOps/s 11.6320 KOps/s $\color{#35bf28}+0.66\%$
test_tc_init_many_fields 35.8220μs 14.7798μs 67.6601 KOps/s 67.6063 KOps/s $\color{#35bf28}+0.08\%$
test_tc_first_layer_tensor 71.8240μs 1.8175μs 550.2044 KOps/s 554.5551 KOps/s $\color{#d91a1a}-0.78\%$
test_tc_first_layer_tensor_only 1.8356μs 0.3986μs 2.5087 MOps/s 2.5371 MOps/s $\color{#d91a1a}-1.12\%$
test_tc_first_layer_tensor_set 37.4420μs 3.9977μs 250.1419 KOps/s 256.1252 KOps/s $\color{#d91a1a}-2.34\%$
test_tc_first_layer_tensor_only_set 24.1710μs 3.2836μs 304.5430 KOps/s 307.7521 KOps/s $\color{#d91a1a}-1.04\%$
test_tc_first_layer_nontensor 35.4820μs 6.7874μs 147.3321 KOps/s 162.7939 KOps/s $\textbf{\color{#d91a1a}-9.50\%}$
test_tc_second_layer_tensor 23.3220μs 4.4676μs 223.8363 KOps/s 226.8811 KOps/s $\color{#d91a1a}-1.34\%$
test_tc_second_layer_nontensor 36.3120μs 9.3847μs 106.5569 KOps/s 116.0893 KOps/s $\textbf{\color{#d91a1a}-8.21\%}$
test_unbind 0.2782s 14.0301ms 71.2752 Ops/s 63.6794 Ops/s $\textbf{\color{#35bf28}+11.93\%}$
test_full_like 4.5897ms 4.3930ms 227.6325 Ops/s 229.3281 Ops/s $\color{#d91a1a}-0.74\%$
test_zeros_like 4.4778ms 4.3368ms 230.5831 Ops/s 232.0673 Ops/s $\color{#d91a1a}-0.64\%$
test_ones_like 4.4436ms 4.3114ms 231.9431 Ops/s 231.3965 Ops/s $\color{#35bf28}+0.24\%$
test_clone 6.8982ms 6.5169ms 153.4480 Ops/s 153.4007 Ops/s $\color{#35bf28}+0.03\%$
test_squeeze 92.8050μs 14.6153μs 68.4215 KOps/s 70.6389 KOps/s $\color{#d91a1a}-3.14\%$
test_unsqueeze 0.1644ms 0.1133ms 8.8278 KOps/s 8.5146 KOps/s $\color{#35bf28}+3.68\%$
test_split 0.2500ms 0.1802ms 5.5489 KOps/s 5.1475 KOps/s $\textbf{\color{#35bf28}+7.80\%}$
test_permute 0.2705ms 0.2072ms 4.8252 KOps/s 4.5519 KOps/s $\textbf{\color{#35bf28}+6.00\%}$
test_stack 51.8685ms 51.4276ms 19.4448 Ops/s 19.5100 Ops/s $\color{#d91a1a}-0.33\%$
test_cat 51.8051ms 51.4693ms 19.4291 Ops/s 19.4766 Ops/s $\color{#d91a1a}-0.24\%$
test_sequential_tensordict 0.9533ms 0.5660ms 1.7668 KOps/s 1.7366 KOps/s $\color{#35bf28}+1.74\%$
test_sequential_graph_module 0.2359ms 0.1199ms 8.3373 KOps/s 8.3193 KOps/s $\color{#35bf28}+0.22\%$
test_nested_tensordict 1.0921ms 0.8145ms 1.2278 KOps/s 1.2133 KOps/s $\color{#35bf28}+1.20\%$
test_nested_graph_module 0.1891ms 0.1324ms 7.5529 KOps/s 7.4548 KOps/s $\color{#35bf28}+1.32\%$

@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 May 24, 2026
@vmoens
Update
b45637e 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request May 24, 2026
@vmoens @cursoragent
[Compile] Make lock_() and other _as_context_manager ops compile-…
cf4c4e6 
…safe

The ``_as_context_manager`` decorator (used by ``lock_``/``unlock_``,
``to_module``, ``transpose``, ``flatten``, etc.) wrote a
``weakref.ref(self)`` into ``_last_op`` so the corresponding ``__exit__``
could call the inverse op. Under ``torch.compile`` that weakref becomes
a ``WeakRefVariable`` that Dynamo cannot reconstruct when the holding
TD flows out of the compiled region, which forced callers to defensively
call ``clear_refs_for_compile_()`` (and previously emitted a strong
warning around ``lock_`` itself).

Use a strong-ref closure (``_strong_ref(obj)`` returning ``lambda: obj``)
in place of ``weakref.ref`` when ``is_compiling()``. The closure is a
plain Python callable Dynamo can handle, and it's safe because the
compiled graph is transient: the closure is dropped as soon as the
outer call ends, so no real lifecycle cycle can form.

Also:

- Add ``_last_op = None`` as a class-level default on ``TensorDictBase``
  so ``__enter__`` can always read it (the compile path no longer writes
  it on every decorated call).
- Update the user-facing ``lock_`` warning to describe the real
  remaining caveat under compile: the parent/child unlock invariant is
  not enforced, the previous weakref-reconstruction failure is gone.

Adds a regression test asserting that a TD locked inside a compiled
function does not leak a real ``weakref.ref`` into ``_last_op``, and
that the strong-ref closure resolves back to the same TD.

Authored with Claude.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: 3f036ad
Pull-Request: #1714
@vmoens
vmoens merged commit b45637e into gh/vmoens/94/base May 24, 2026
28 of 44 checks passed
@vmoens
vmoens deleted the gh/vmoens/94/head branch May 24, 2026 14:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Compile torch.compile related Test

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens