Skip to content

[Compile] Locked-fast-path for clone(recurse=False) via cached schema - #1715

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

[Compile] Locked-fast-path for clone(recurse=False) via cached schema#1715
vmoens merged 2 commits into
gh/vmoens/95/basefrom
gh/vmoens/95/head

Conversation

@vmoens

@vmoens vmoens commented May 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

A locked TensorDict already guarantees that its top-level key set
is fixed. Cache that key set in a _LockedSchema namedtuple when
lock_() is called and consume it in the hot path so Dynamo can
guard on an immutable tuple instead of walking the backing dict.

Concretely, under torch.compile clone(recurse=False) on a
locked TD switches from

result._tensordict.update(self._tensordict)

(which forces a DICT_LENGTH / dict-keys guard on _tensordict)
to a per-key copy driven by the cached schema.keys tuple. Cheap
LENGTH_CHECK / EQUALS_MATCH guards on the immutable schema
replace the dict-internals guard, and the schema can be shared
across locked TDs that lock to the same key set.

Plumbing:

  • _LockedSchema namedtuple added in tensordict.utils.
  • TensorDictBase._locked_schema class default + new
    _build_lock_schema() hook called from _propagate_lock in
    eager mode (skipped under compile to avoid the very guard we want
    to avoid building).
  • _propagate_unlock drops the schema so it cannot go stale once
    keys can be added/removed again.
  • TensorDict._build_lock_schema snapshots tuple(self._tensordict).
  • TensorDict._clone uses the cached schema under
    is_compiling() and self._is_locked.

Adds guard-count tests that exercise the locked clone fast path and
verify the schema is cleared on unlock and rebuilt on relock.

The locked-fast-path infrastructure is intentionally narrow in this
commit: only clone(recurse=False) consumes the schema. Extending
it to _set_str / _has_names is left for follow-ups so the
schema-cache invariants can be evaluated under real workloads first.

Authored with Claude.

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

@vmoens
Update
9654228 
[ghstack-poisoned]
This was referenced May 24, 2026
Merged
Merged
@github-actions github-actions Bot added Test Compile torch.compile related labels May 24, 2026
@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 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}16$. Worsened: $\large\color{#d91a1a}14$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_collector_preallocated_write[stack_out-cpu-32-2] 0.2261ms 0.1727ms 5.7918 KOps/s 6.1615 KOps/s $\textbf{\color{#d91a1a}-6.00\%}$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.3667ms 0.2936ms 3.4063 KOps/s 3.2957 KOps/s $\color{#35bf28}+3.35\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8445ms 0.7470ms 1.3386 KOps/s 1.3205 KOps/s $\color{#35bf28}+1.37\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.3835ms 0.3465ms 2.8856 KOps/s 2.8409 KOps/s $\color{#35bf28}+1.57\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.6066ms 1.0980ms 910.7141 Ops/s 892.1204 Ops/s $\color{#35bf28}+2.08\%$
test_collector_preallocated_write[stack_out-cpu-4096-32] 7.6013ms 7.2145ms 138.6107 Ops/s 137.4108 Ops/s $\color{#35bf28}+0.87\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.3031ms 0.2443ms 4.0933 KOps/s 3.9369 KOps/s $\color{#35bf28}+3.97\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.4633ms 0.3760ms 2.6598 KOps/s 2.6585 KOps/s $\color{#35bf28}+0.05\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.9184ms 0.7905ms 1.2650 KOps/s 1.2472 KOps/s $\color{#35bf28}+1.42\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.3262ms 0.2419ms 4.1344 KOps/s 3.8846 KOps/s $\textbf{\color{#35bf28}+6.43\%}$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.4481ms 0.3623ms 2.7602 KOps/s 2.7538 KOps/s $\color{#35bf28}+0.23\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 0.9781ms 0.8564ms 1.1676 KOps/s 1.1661 KOps/s $\color{#35bf28}+0.13\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.2844ms 0.2127ms 4.7010 KOps/s 4.7658 KOps/s $\color{#d91a1a}-1.36\%$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 0.9048ms 0.8223ms 1.2160 KOps/s 1.2080 KOps/s $\color{#35bf28}+0.67\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.3855ms 3.2758ms 305.2696 Ops/s 304.2698 Ops/s $\color{#35bf28}+0.33\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.7382ms 0.3584ms 2.7899 KOps/s 2.7506 KOps/s $\color{#35bf28}+1.43\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 1.8124ms 1.6558ms 603.9304 Ops/s 617.2297 Ops/s $\color{#d91a1a}-2.15\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 10.5353ms 10.2105ms 97.9388 Ops/s 99.0790 Ops/s $\color{#d91a1a}-1.15\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.4840ms 0.4042ms 2.4740 KOps/s 2.5765 KOps/s $\color{#d91a1a}-3.98\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.6957ms 1.5108ms 661.9076 Ops/s 647.0509 Ops/s $\color{#35bf28}+2.30\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 6.2663ms 5.8903ms 169.7703 Ops/s 163.6131 Ops/s $\color{#35bf28}+3.76\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.4763ms 0.4158ms 2.4053 KOps/s 2.4075 KOps/s $\color{#d91a1a}-0.09\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.7715ms 1.5314ms 653.0074 Ops/s 646.0890 Ops/s $\color{#35bf28}+1.07\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 6.3083ms 6.0338ms 165.7323 Ops/s 165.8854 Ops/s $\color{#d91a1a}-0.09\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.2301ms 0.1702ms 5.8749 KOps/s 5.8577 KOps/s $\color{#35bf28}+0.29\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 0.7473ms 0.6621ms 1.5102 KOps/s 1.4942 KOps/s $\color{#35bf28}+1.07\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 2.8198ms 2.6463ms 377.8929 Ops/s 376.6556 Ops/s $\color{#35bf28}+0.33\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.3884ms 0.3180ms 3.1448 KOps/s 3.1505 KOps/s $\color{#d91a1a}-0.18\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 1.9141ms 1.5038ms 664.9644 Ops/s 714.2335 Ops/s $\textbf{\color{#d91a1a}-6.90\%}$
test_collector_preallocated_write[copy_at-cpu-4096-32] 9.7737ms 9.5067ms 105.1888 Ops/s 107.2376 Ops/s $\color{#d91a1a}-1.91\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.3937ms 0.3274ms 3.0544 KOps/s 3.1269 KOps/s $\color{#d91a1a}-2.32\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.3112ms 1.2219ms 818.3909 Ops/s 838.8317 Ops/s $\color{#d91a1a}-2.44\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 4.9444ms 4.8042ms 208.1533 Ops/s 214.8250 Ops/s $\color{#d91a1a}-3.11\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.3831ms 0.3276ms 3.0523 KOps/s 3.1143 KOps/s $\color{#d91a1a}-1.99\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.3319ms 1.2305ms 812.6785 Ops/s 832.8223 Ops/s $\color{#d91a1a}-2.42\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 4.9867ms 4.8049ms 208.1197 Ops/s 212.7032 Ops/s $\color{#d91a1a}-2.15\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1704ms 0.1197ms 8.3538 KOps/s 8.5615 KOps/s $\color{#d91a1a}-2.43\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5527ms 0.4722ms 2.1175 KOps/s 2.1743 KOps/s $\color{#d91a1a}-2.61\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 1.9886ms 1.8749ms 533.3758 Ops/s 549.4613 Ops/s $\color{#d91a1a}-2.93\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.3444ms 0.2667ms 3.7502 KOps/s 3.9586 KOps/s $\textbf{\color{#d91a1a}-5.26\%}$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.6507ms 1.3016ms 768.3040 Ops/s 810.6802 Ops/s $\textbf{\color{#d91a1a}-5.23\%}$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 8.9737ms 8.4428ms 118.4441 Ops/s 118.2227 Ops/s $\color{#35bf28}+0.19\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.3750ms 0.2832ms 3.5314 KOps/s 3.6841 KOps/s $\color{#d91a1a}-4.14\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.1348ms 1.0576ms 945.5383 Ops/s 980.7279 Ops/s $\color{#d91a1a}-3.59\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.2797ms 4.1576ms 240.5236 Ops/s 254.5188 Ops/s $\textbf{\color{#d91a1a}-5.50\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.3711ms 0.2830ms 3.5334 KOps/s 3.6626 KOps/s $\color{#d91a1a}-3.53\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.1486ms 1.0625ms 941.2200 Ops/s 980.9979 Ops/s $\color{#d91a1a}-4.05\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.3393ms 4.1784ms 239.3264 Ops/s 251.4630 Ops/s $\color{#d91a1a}-4.83\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 16.0345ms 15.5624ms 64.2573 Ops/s 63.8771 Ops/s $\color{#35bf28}+0.60\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 31.8215ms 31.5901ms 31.6555 Ops/s 31.7936 Ops/s $\color{#d91a1a}-0.43\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 83.0673ms 82.8464ms 12.0705 Ops/s 12.1172 Ops/s $\color{#d91a1a}-0.38\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2071s 0.2028s 4.9299 Ops/s 5.1861 Ops/s $\color{#d91a1a}-4.94\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 13.4395ms 13.1578ms 76.0004 Ops/s 76.1540 Ops/s $\color{#d91a1a}-0.20\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 26.1717ms 25.5900ms 39.0778 Ops/s 39.0027 Ops/s $\color{#35bf28}+0.19\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 13.6407ms 13.3518ms 74.8964 Ops/s 75.3841 Ops/s $\color{#d91a1a}-0.65\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 26.1062ms 25.8725ms 38.6511 Ops/s 38.2259 Ops/s $\color{#35bf28}+1.11\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 52.0112ms 51.6899ms 19.3461 Ops/s 19.4418 Ops/s $\color{#d91a1a}-0.49\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1049s 0.1042s 9.5960 Ops/s 9.6811 Ops/s $\color{#d91a1a}-0.88\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1455s 0.1440s 6.9447 Ops/s 7.0464 Ops/s $\color{#d91a1a}-1.44\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3239s 0.3212s 3.1134 Ops/s 3.2234 Ops/s $\color{#d91a1a}-3.41\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 79.0878ms 74.3361ms 13.4524 Ops/s 13.6237 Ops/s $\color{#d91a1a}-1.26\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1576s 0.1542s 6.4832 Ops/s 6.8181 Ops/s $\color{#d91a1a}-4.91\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 74.9816ms 74.4898ms 13.4247 Ops/s 13.4233 Ops/s $\color{#35bf28}+0.01\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1516s 0.1513s 6.6111 Ops/s 6.7439 Ops/s $\color{#d91a1a}-1.97\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.4142ms 29.9931ms 33.3410 Ops/s 33.3657 Ops/s $\color{#d91a1a}-0.07\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 61.4870ms 61.0363ms 16.3837 Ops/s 16.4374 Ops/s $\color{#d91a1a}-0.33\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1162s 0.1158s 8.6340 Ops/s 8.6816 Ops/s $\color{#d91a1a}-0.55\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2602s 0.2592s 3.8573 Ops/s 3.8927 Ops/s $\color{#d91a1a}-0.91\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 44.0316ms 43.5263ms 22.9746 Ops/s 23.5261 Ops/s $\color{#d91a1a}-2.34\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 88.5068ms 87.5218ms 11.4257 Ops/s 11.7254 Ops/s $\color{#d91a1a}-2.56\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 44.1449ms 43.8113ms 22.8251 Ops/s 23.4000 Ops/s $\color{#d91a1a}-2.46\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 89.0619ms 87.7792ms 11.3922 Ops/s 11.6581 Ops/s $\color{#d91a1a}-2.28\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 20.8052ms 20.4936ms 48.7958 Ops/s 50.1870 Ops/s $\color{#d91a1a}-2.77\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 42.4652ms 42.0590ms 23.7762 Ops/s 24.6065 Ops/s $\color{#d91a1a}-3.37\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1058s 0.1054s 9.4843 Ops/s 9.6175 Ops/s $\color{#d91a1a}-1.39\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2400s 0.2399s 4.1687 Ops/s 4.2972 Ops/s $\color{#d91a1a}-2.99\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 38.2682ms 37.7234ms 26.5087 Ops/s 27.6764 Ops/s $\color{#d91a1a}-4.22\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 76.6578ms 75.8840ms 13.1780 Ops/s 13.7297 Ops/s $\color{#d91a1a}-4.02\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 38.7656ms 38.1989ms 26.1787 Ops/s 27.3201 Ops/s $\color{#d91a1a}-4.18\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 77.9144ms 76.8457ms 13.0131 Ops/s 13.1666 Ops/s $\color{#d91a1a}-1.17\%$
test_plain_set_nested 52.7710μs 15.1775μs 65.8869 KOps/s 66.9654 KOps/s $\color{#d91a1a}-1.61\%$
test_plain_set_stack_nested 37.6510μs 15.0118μs 66.6142 KOps/s 67.5086 KOps/s $\color{#d91a1a}-1.32\%$
test_plain_set_nested_inplace 52.6610μs 16.3812μs 61.0455 KOps/s 61.4844 KOps/s $\color{#d91a1a}-0.71\%$
test_plain_set_stack_nested_inplace 38.0910μs 16.1276μs 62.0055 KOps/s 62.5787 KOps/s $\color{#d91a1a}-0.92\%$
test_items 44.1110μs 5.4790μs 182.5136 KOps/s 181.0469 KOps/s $\color{#35bf28}+0.81\%$
test_items_nested 0.5065ms 0.4469ms 2.2378 KOps/s 2.1788 KOps/s $\color{#35bf28}+2.71\%$
test_items_nested_locked 0.5681ms 0.4496ms 2.2240 KOps/s 2.1911 KOps/s $\color{#35bf28}+1.50\%$
test_items_nested_leaf 0.1366ms 92.7379μs 10.7831 KOps/s 10.7423 KOps/s $\color{#35bf28}+0.38\%$
test_items_stack_nested 0.5414ms 0.4474ms 2.2350 KOps/s 2.2015 KOps/s $\color{#35bf28}+1.52\%$
test_items_stack_nested_leaf 0.1335ms 93.0752μs 10.7440 KOps/s 10.6954 KOps/s $\color{#35bf28}+0.45\%$
test_items_stack_nested_locked 0.5262ms 0.4506ms 2.2194 KOps/s 2.1860 KOps/s $\color{#35bf28}+1.53\%$
test_keys 34.8300μs 4.3245μs 231.2404 KOps/s 238.9551 KOps/s $\color{#d91a1a}-3.23\%$
test_keys_nested 0.2162ms 0.1297ms 7.7120 KOps/s 7.6892 KOps/s $\color{#35bf28}+0.30\%$
test_keys_nested_locked 0.7750ms 0.1382ms 7.2369 KOps/s 7.2123 KOps/s $\color{#35bf28}+0.34\%$
test_keys_nested_leaf 0.1713ms 0.1202ms 8.3197 KOps/s 8.2722 KOps/s $\color{#35bf28}+0.57\%$
test_keys_stack_nested 0.1780ms 0.1293ms 7.7337 KOps/s 7.6785 KOps/s $\color{#35bf28}+0.72\%$
test_keys_stack_nested_leaf 0.1512ms 0.1188ms 8.4165 KOps/s 8.3525 KOps/s $\color{#35bf28}+0.77\%$
test_keys_stack_nested_locked 0.1725ms 0.1366ms 7.3216 KOps/s 7.2620 KOps/s $\color{#35bf28}+0.82\%$
test_values 5.9740μs 0.9992μs 1.0008 MOps/s 996.3455 KOps/s $\color{#35bf28}+0.45\%$
test_values_nested 0.1006ms 51.8645μs 19.2810 KOps/s 19.0621 KOps/s $\color{#35bf28}+1.15\%$
test_values_nested_locked 89.1510μs 55.1973μs 18.1168 KOps/s 18.1453 KOps/s $\color{#d91a1a}-0.16\%$
test_values_nested_leaf 87.6910μs 58.9465μs 16.9645 KOps/s 16.6667 KOps/s $\color{#35bf28}+1.79\%$
test_values_stack_nested 90.1910μs 51.5599μs 19.3949 KOps/s 19.1585 KOps/s $\color{#35bf28}+1.23\%$
test_values_stack_nested_leaf 95.5110μs 58.7937μs 17.0086 KOps/s 16.7986 KOps/s $\color{#35bf28}+1.25\%$
test_values_stack_nested_locked 88.9610μs 54.8403μs 18.2348 KOps/s 18.0238 KOps/s $\color{#35bf28}+1.17\%$
test_membership 4.6018μs 0.7952μs 1.2575 MOps/s 1.2481 MOps/s $\color{#35bf28}+0.75\%$
test_membership_nested 20.6800μs 2.7531μs 363.2230 KOps/s 354.7330 KOps/s $\color{#35bf28}+2.39\%$
test_membership_nested_leaf 19.7955μs 2.6734μs 374.0558 KOps/s 356.7114 KOps/s $\color{#35bf28}+4.86\%$
test_membership_stacked_nested 31.0000μs 2.7252μs 366.9403 KOps/s 365.2850 KOps/s $\color{#35bf28}+0.45\%$
test_membership_stacked_nested_leaf 30.0500μs 2.7362μs 365.4713 KOps/s 363.4386 KOps/s $\color{#35bf28}+0.56\%$
test_membership_nested_last 43.0110μs 4.2047μs 237.8285 KOps/s 238.0931 KOps/s $\color{#d91a1a}-0.11\%$
test_membership_nested_leaf_last 35.3200μs 4.1884μs 238.7558 KOps/s 236.8767 KOps/s $\color{#35bf28}+0.79\%$
test_membership_stacked_nested_last 31.2110μs 4.1536μs 240.7530 KOps/s 240.2129 KOps/s $\color{#35bf28}+0.22\%$
test_membership_stacked_nested_leaf_last 33.4410μs 4.1138μs 243.0862 KOps/s 244.4308 KOps/s $\color{#d91a1a}-0.55\%$
test_nested_getleaf 57.6610μs 21.1644μs 47.2491 KOps/s 48.1300 KOps/s $\color{#d91a1a}-1.83\%$
test_nested_get 57.4910μs 19.8392μs 50.4052 KOps/s 51.4582 KOps/s $\color{#d91a1a}-2.05\%$
test_stacked_getleaf 44.0510μs 20.4405μs 48.9226 KOps/s 48.2617 KOps/s $\color{#35bf28}+1.37\%$
test_stacked_get 58.3510μs 19.7656μs 50.5931 KOps/s 51.1429 KOps/s $\color{#d91a1a}-1.08\%$
test_nested_getitemleaf 50.2400μs 21.2389μs 47.0835 KOps/s 47.1797 KOps/s $\color{#d91a1a}-0.20\%$
test_nested_getitem 71.0510μs 20.0879μs 49.7813 KOps/s 49.3356 KOps/s $\color{#35bf28}+0.90\%$
test_stacked_getitemleaf 53.7910μs 20.9208μs 47.7992 KOps/s 47.3269 KOps/s $\color{#35bf28}+1.00\%$
test_stacked_getitem 57.2210μs 20.0867μs 49.7841 KOps/s 49.6477 KOps/s $\color{#35bf28}+0.27\%$
test_lock_nested 0.5725ms 0.4792ms 2.0867 KOps/s 2.1334 KOps/s $\color{#d91a1a}-2.19\%$
test_lock_stack_nested 0.5525ms 0.4886ms 2.0465 KOps/s 2.1461 KOps/s $\color{#d91a1a}-4.64\%$
test_unlock_nested 0.4548ms 0.3863ms 2.5885 KOps/s 2.6461 KOps/s $\color{#d91a1a}-2.18\%$
test_unlock_stack_nested 0.4415ms 0.3882ms 2.5761 KOps/s 2.6238 KOps/s $\color{#d91a1a}-1.82\%$
test_flatten_speed 0.1939ms 0.1158ms 8.6354 KOps/s 8.6056 KOps/s $\color{#35bf28}+0.35\%$
test_unflatten_speed 0.6566ms 0.5488ms 1.8222 KOps/s 1.8237 KOps/s $\color{#d91a1a}-0.08\%$
test_common_ops 0.1276s 1.6062ms 622.5713 Ops/s 430.1310 Ops/s $\textbf{\color{#35bf28}+44.74\%}$
test_creation 0.1131ms 2.9745μs 336.1963 KOps/s 344.8517 KOps/s $\color{#d91a1a}-2.51\%$
test_creation_empty 27.1300μs 6.5743μs 152.1063 KOps/s 155.2568 KOps/s $\color{#d91a1a}-2.03\%$
test_creation_nested_1 39.6910μs 10.8273μs 92.3587 KOps/s 92.2445 KOps/s $\color{#35bf28}+0.12\%$
test_creation_nested_2 39.8100μs 12.8926μs 77.5641 KOps/s 79.0038 KOps/s $\color{#d91a1a}-1.82\%$
test_creation_many_keys[10] 44.0910μs 19.6873μs 50.7941 KOps/s 51.8094 KOps/s $\color{#d91a1a}-1.96\%$
test_creation_many_keys[50] 0.1245ms 84.5300μs 11.8301 KOps/s 12.0419 KOps/s $\color{#d91a1a}-1.76\%$
test_creation_many_keys[100] 0.2297ms 0.1667ms 5.9989 KOps/s 6.1425 KOps/s $\color{#d91a1a}-2.34\%$
test_creation_nested_many_keys[10] 94.5320μs 42.3384μs 23.6192 KOps/s 23.7628 KOps/s $\color{#d91a1a}-0.60\%$
test_creation_nested_many_keys[50] 0.2147ms 0.1719ms 5.8163 KOps/s 5.8966 KOps/s $\color{#d91a1a}-1.36\%$
test_clone 48.1610μs 13.0172μs 76.8214 KOps/s 77.1224 KOps/s $\color{#d91a1a}-0.39\%$
test_getitem[int] 1.5496ms 14.5943μs 68.5198 KOps/s 68.8275 KOps/s $\color{#d91a1a}-0.45\%$
test_getitem[slice_int] 0.1368ms 23.9614μs 41.7338 KOps/s 42.7525 KOps/s $\color{#d91a1a}-2.38\%$
test_getitem[range] 0.1737ms 62.4287μs 16.0183 KOps/s 15.7628 KOps/s $\color{#35bf28}+1.62\%$
test_getitem[tuple] 0.1725ms 23.7516μs 42.1024 KOps/s 42.0176 KOps/s $\color{#35bf28}+0.20\%$
test_getitem[list] 0.1756ms 57.2211μs 17.4761 KOps/s 17.2299 KOps/s $\color{#35bf28}+1.43\%$
test_setitem_dim[int] 45.6810μs 24.9943μs 40.0091 KOps/s 39.6076 KOps/s $\color{#35bf28}+1.01\%$
test_setitem_dim[slice_int] 79.2010μs 42.5973μs 23.4757 KOps/s 24.1866 KOps/s $\color{#d91a1a}-2.94\%$
test_setitem_dim[range] 0.1385ms 94.1935μs 10.6164 KOps/s 10.5354 KOps/s $\color{#35bf28}+0.77\%$
test_setitem_dim[tuple] 60.8310μs 39.5028μs 25.3146 KOps/s 25.5527 KOps/s $\color{#d91a1a}-0.93\%$
test_setitem 49.9510μs 17.5060μs 57.1231 KOps/s 58.7908 KOps/s $\color{#d91a1a}-2.84\%$
test_set 53.5410μs 16.4719μs 60.7093 KOps/s 61.2320 KOps/s $\color{#d91a1a}-0.85\%$
test_set_shared 0.5458ms 0.2103ms 4.7562 KOps/s 4.8899 KOps/s $\color{#d91a1a}-2.73\%$
test_update 0.1980ms 21.7526μs 45.9715 KOps/s 46.7119 KOps/s $\color{#d91a1a}-1.58\%$
test_update_nested 69.2320μs 33.5258μs 29.8278 KOps/s 30.7949 KOps/s $\color{#d91a1a}-3.14\%$
test_update__nested 0.4969ms 32.8384μs 30.4521 KOps/s 30.6595 KOps/s $\color{#d91a1a}-0.68\%$
test_set_nested 0.1164ms 18.4295μs 54.2608 KOps/s 55.6847 KOps/s $\color{#d91a1a}-2.56\%$
test_set_nested_new 57.7410μs 23.3775μs 42.7761 KOps/s 43.5829 KOps/s $\color{#d91a1a}-1.85\%$
test_select 77.5110μs 39.9942μs 25.0036 KOps/s 26.1779 KOps/s $\color{#d91a1a}-4.49\%$
test_select_nested 0.1128ms 69.7406μs 14.3389 KOps/s 14.1523 KOps/s $\color{#35bf28}+1.32\%$
test_exclude_nested 0.1301ms 85.7210μs 11.6658 KOps/s 11.3733 KOps/s $\color{#35bf28}+2.57\%$
test_empty[True] 0.4915ms 0.3923ms 2.5493 KOps/s 2.5455 KOps/s $\color{#35bf28}+0.15\%$
test_empty[False] 9.9103μs 1.3220μs 756.4137 KOps/s 766.8658 KOps/s $\color{#d91a1a}-1.36\%$
test_to 0.1043ms 67.2877μs 14.8616 KOps/s 14.8850 KOps/s $\color{#d91a1a}-0.16\%$
test_to_nonblocking 0.1095ms 67.5063μs 14.8134 KOps/s 14.7351 KOps/s $\color{#35bf28}+0.53\%$
test_unbind_speed 0.3952ms 0.3299ms 3.0313 KOps/s 3.0526 KOps/s $\color{#d91a1a}-0.70\%$
test_unbind_speed_stack0 0.4249ms 0.3261ms 3.0670 KOps/s 3.1204 KOps/s $\color{#d91a1a}-1.71\%$
test_unbind_speed_stack1 0.1151s 0.8417ms 1.1881 KOps/s 1.1919 KOps/s $\color{#d91a1a}-0.31\%$
test_split 1.2030ms 1.1001ms 909.0334 Ops/s 724.0876 Ops/s $\textbf{\color{#35bf28}+25.54\%}$
test_chunk 0.1148s 1.1876ms 842.0502 Ops/s 953.1565 Ops/s $\textbf{\color{#d91a1a}-11.66\%}$
test_to_cpu_blocking 20.9361ms 20.1394ms 49.6540 Ops/s 49.5661 Ops/s $\color{#35bf28}+0.18\%$
test_to_cpu_global_sync 12.0671ms 11.7626ms 85.0154 Ops/s 84.9370 Ops/s $\color{#35bf28}+0.09\%$
test_to_cpu_event_sync 13.3095ms 12.7661ms 78.3324 Ops/s 78.7399 Ops/s $\color{#d91a1a}-0.52\%$
test_to_cpu_default 0.1279s 14.2257ms 70.2955 Ops/s 70.1350 Ops/s $\color{#35bf28}+0.23\%$
test_consolidate[False-None] 0.1439s 4.6392ms 215.5541 Ops/s 251.9081 Ops/s $\textbf{\color{#d91a1a}-14.43\%}$
test_consolidate[default-None] 0.1864s 2.4733ms 404.3153 Ops/s 465.6532 Ops/s $\textbf{\color{#d91a1a}-13.17\%}$
test_consolidate[reduce-overhead-None] 0.2016s 2.3843ms 419.4094 Ops/s 479.7931 Ops/s $\textbf{\color{#d91a1a}-12.59\%}$
test_consolidate_njt[False-None] 8.4351ms 8.1902ms 122.0973 Ops/s 119.2367 Ops/s $\color{#35bf28}+2.40\%$
test_to[False-False-None] 2.2469ms 2.1293ms 469.6467 Ops/s 464.3104 Ops/s $\color{#35bf28}+1.15\%$
test_to[True-False-None] 2.1455ms 1.8646ms 536.3152 Ops/s 546.3696 Ops/s $\color{#d91a1a}-1.84\%$
test_to[within-False-None] 6.3378ms 6.0433ms 165.4719 Ops/s 168.3576 Ops/s $\color{#d91a1a}-1.71\%$
test_to[True-default-None] 10.6205ms 10.4910ms 95.3200 Ops/s 88.1207 Ops/s $\textbf{\color{#35bf28}+8.17\%}$
test_to_njt[False-False-None] 8.4800ms 8.3749ms 119.4051 Ops/s 110.6620 Ops/s $\textbf{\color{#35bf28}+7.90\%}$
test_to_njt[True-False-None] 6.8714ms 6.7045ms 149.1536 Ops/s 137.6867 Ops/s $\textbf{\color{#35bf28}+8.33\%}$
test_to_njt[within-False-None] 15.1808ms 14.8512ms 67.3348 Ops/s 63.2358 Ops/s $\textbf{\color{#35bf28}+6.48\%}$
test_creation[device0] 0.4288ms 0.1128ms 8.8655 KOps/s 8.7312 KOps/s $\color{#35bf28}+1.54\%$
test_creation_from_tensor 0.4223ms 0.1106ms 9.0420 KOps/s 8.7845 KOps/s $\color{#35bf28}+2.93\%$
test_add_one[memmap_tensor0] 0.2893ms 6.6286μs 150.8611 KOps/s 149.4581 KOps/s $\color{#35bf28}+0.94\%$
test_contiguous[memmap_tensor0] 13.8900μs 0.5993μs 1.6687 MOps/s 2.1410 MOps/s $\textbf{\color{#d91a1a}-22.06\%}$
test_stack[memmap_tensor0] 50.6310μs 4.0808μs 245.0525 KOps/s 239.9316 KOps/s $\color{#35bf28}+2.13\%$
test_memmaptd_index 1.3028ms 0.2656ms 3.7646 KOps/s 3.7438 KOps/s $\color{#35bf28}+0.56\%$
test_memmaptd_index_astensor 0.5152ms 0.3645ms 2.7436 KOps/s 2.7263 KOps/s $\color{#35bf28}+0.64\%$
test_memmaptd_index_op 0.9324ms 0.6184ms 1.6170 KOps/s 1.5892 KOps/s $\color{#35bf28}+1.75\%$
test_serialize_model 0.1389s 0.1354s 7.3867 Ops/s 7.4177 Ops/s $\color{#d91a1a}-0.42\%$
test_serialize_model_pickle 2.0589s 1.3281s 0.7530 Ops/s 0.8265 Ops/s $\textbf{\color{#d91a1a}-8.90\%}$
test_serialize_weights 0.1359s 0.1347s 7.4216 Ops/s 7.4553 Ops/s $\color{#d91a1a}-0.45\%$
test_serialize_weights_returnearly 0.4610s 87.3403ms 11.4495 Ops/s 6.5355 Ops/s $\textbf{\color{#35bf28}+75.19\%}$
test_serialize_weights_pickle 1.4010s 1.2227s 0.8178 Ops/s 0.8222 Ops/s $\color{#d91a1a}-0.53\%$
test_reshape_pytree 0.2036ms 31.8009μs 31.4456 KOps/s 31.7133 KOps/s $\color{#d91a1a}-0.84\%$
test_reshape_td 94.7910μs 43.1501μs 23.1749 KOps/s 23.1630 KOps/s $\color{#35bf28}+0.05\%$
test_view_pytree 0.2139ms 30.8714μs 32.3925 KOps/s 30.6314 KOps/s $\textbf{\color{#35bf28}+5.75\%}$
test_view_td 0.1063ms 53.3272μs 18.7522 KOps/s 18.1478 KOps/s $\color{#35bf28}+3.33\%$
test_unbind_pytree 0.2425ms 36.1680μs 27.6488 KOps/s 26.5372 KOps/s $\color{#35bf28}+4.19\%$
test_unbind_td 0.2113ms 49.2796μs 20.2924 KOps/s 20.4128 KOps/s $\color{#d91a1a}-0.59\%$
test_split_pytree 0.2470ms 41.2413μs 24.2475 KOps/s 23.0527 KOps/s $\textbf{\color{#35bf28}+5.18\%}$
test_split_td 0.2057ms 62.5965μs 15.9753 KOps/s 15.4385 KOps/s $\color{#35bf28}+3.48\%$
test_add_pytree 0.2285ms 41.5500μs 24.0674 KOps/s 22.8973 KOps/s $\textbf{\color{#35bf28}+5.11\%}$
test_add_td 0.1048ms 57.8570μs 17.2840 KOps/s 16.6428 KOps/s $\color{#35bf28}+3.85\%$
test_compile_add_one_nested[tensordict-compile] 0.3605ms 0.2351ms 4.2540 KOps/s 4.3225 KOps/s $\color{#d91a1a}-1.58\%$
test_compile_add_one_nested[tensordict-eager] 0.2886ms 0.1948ms 5.1333 KOps/s 5.1558 KOps/s $\color{#d91a1a}-0.44\%$
test_compile_add_one_nested[pytree-compile] 0.4043ms 0.1399ms 7.1501 KOps/s 7.3231 KOps/s $\color{#d91a1a}-2.36\%$
test_compile_add_one_nested[pytree-eager] 0.4236ms 0.1757ms 5.6902 KOps/s 5.6995 KOps/s $\color{#d91a1a}-0.16\%$
test_compile_copy_nested[tensordict-compile] 0.1329ms 15.2414μs 65.6106 KOps/s 67.8421 KOps/s $\color{#d91a1a}-3.29\%$
test_compile_copy_nested[tensordict-eager] 0.3059ms 50.6136μs 19.7575 KOps/s 19.3164 KOps/s $\color{#35bf28}+2.28\%$
test_compile_copy_nested[pytree-compile] 0.1592ms 15.5777μs 64.1945 KOps/s 65.7368 KOps/s $\color{#d91a1a}-2.35\%$
test_compile_copy_nested[pytree-eager] 0.3729ms 64.1877μs 15.5793 KOps/s 15.4464 KOps/s $\color{#35bf28}+0.86\%$
test_compile_add_one_flat[tensordict-compile] 0.3751ms 0.2555ms 3.9144 KOps/s 3.9539 KOps/s $\color{#d91a1a}-1.00\%$
test_compile_add_one_flat[tensordict-eager] 0.3565ms 0.2773ms 3.6058 KOps/s 3.5997 KOps/s $\color{#35bf28}+0.17\%$
test_compile_add_one_flat[tensorclass-compile] 0.3106ms 0.1864ms 5.3659 KOps/s 5.1270 KOps/s $\color{#35bf28}+4.66\%$
test_compile_add_one_flat[tensorclass-eager] 0.1209ms 75.6153μs 13.2248 KOps/s 13.2466 KOps/s $\color{#d91a1a}-0.16\%$
test_compile_add_one_flat[pytree-compile] 0.4235ms 0.1901ms 5.2606 KOps/s 5.0707 KOps/s $\color{#35bf28}+3.75\%$
test_compile_add_one_flat[pytree-eager] 0.8225ms 0.5287ms 1.8915 KOps/s 1.9600 KOps/s $\color{#d91a1a}-3.50\%$
test_compile_add_self_flat[tensordict-eager] 0.4078ms 0.3226ms 3.0998 KOps/s 3.0442 KOps/s $\color{#35bf28}+1.83\%$
test_compile_add_self_flat[tensordict-compile] 0.3125ms 0.2470ms 4.0479 KOps/s 3.7926 KOps/s $\textbf{\color{#35bf28}+6.73\%}$
test_compile_add_self_flat[tensorclass-eager] 0.1327ms 90.5872μs 11.0391 KOps/s 11.0058 KOps/s $\color{#35bf28}+0.30\%$
test_compile_add_self_flat[tensorclass-compile] 0.3068ms 0.1848ms 5.4116 KOps/s 5.1310 KOps/s $\textbf{\color{#35bf28}+5.47\%}$
test_compile_add_self_flat[pytree-eager] 0.6659ms 0.4333ms 2.3078 KOps/s 2.3407 KOps/s $\color{#d91a1a}-1.41\%$
test_compile_add_self_flat[pytree-compile] 0.3608ms 0.1997ms 5.0068 KOps/s 5.0729 KOps/s $\color{#d91a1a}-1.30\%$
test_compile_copy_flat[tensordict-compile] 58.9610μs 19.1150μs 52.3150 KOps/s 51.7462 KOps/s $\color{#35bf28}+1.10\%$
test_compile_copy_flat[tensordict-eager] 70.6610μs 39.9183μs 25.0512 KOps/s 24.5696 KOps/s $\color{#35bf28}+1.96\%$
test_compile_copy_flat[pytree-compile] 0.2525ms 20.8602μs 47.9381 KOps/s 47.0191 KOps/s $\color{#35bf28}+1.95\%$
test_compile_copy_flat[pytree-eager] 0.3309ms 50.9300μs 19.6348 KOps/s 19.5254 KOps/s $\color{#35bf28}+0.56\%$
test_compile_assign_and_add[tensordict-compile] 2.1517ms 0.2394ms 4.1779 KOps/s 3.9950 KOps/s $\color{#35bf28}+4.58\%$
test_compile_assign_and_add[tensordict-eager] 3.4054ms 3.2408ms 308.5679 Ops/s 309.8357 Ops/s $\color{#d91a1a}-0.41\%$
test_compile_assign_and_add[pytree-compile] 2.0681ms 0.1780ms 5.6191 KOps/s 5.5143 KOps/s $\color{#35bf28}+1.90\%$
test_compile_assign_and_add[pytree-eager] 2.9632ms 2.7868ms 358.8345 Ops/s 367.7813 Ops/s $\color{#d91a1a}-2.43\%$
test_compile_indexing[tensor-tensordict-compile] 0.2890ms 0.1343ms 7.4436 KOps/s 7.3914 KOps/s $\color{#35bf28}+0.71\%$
test_compile_indexing[tensor-tensordict-eager] 0.3484ms 77.0223μs 12.9832 KOps/s 13.0268 KOps/s $\color{#d91a1a}-0.33\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2323ms 0.1224ms 8.1706 KOps/s 8.2029 KOps/s $\color{#d91a1a}-0.39\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2751ms 47.0293μs 21.2633 KOps/s 21.9078 KOps/s $\color{#d91a1a}-2.94\%$
test_compile_indexing[tensor-pytree-compile] 0.1665ms 0.1203ms 8.3101 KOps/s 8.0622 KOps/s $\color{#35bf28}+3.08\%$
test_compile_indexing[tensor-pytree-eager] 0.2579ms 47.7889μs 20.9254 KOps/s 21.6793 KOps/s $\color{#d91a1a}-3.48\%$
test_compile_indexing[slice-tensordict-compile] 0.1863ms 86.8623μs 11.5125 KOps/s 11.4652 KOps/s $\color{#35bf28}+0.41\%$
test_compile_indexing[slice-tensordict-eager] 0.2018ms 29.7980μs 33.5593 KOps/s 35.4695 KOps/s $\textbf{\color{#d91a1a}-5.39\%}$
test_compile_indexing[slice-tensorclass-compile] 0.1869ms 71.6335μs 13.9599 KOps/s 13.9306 KOps/s $\color{#35bf28}+0.21\%$
test_compile_indexing[slice-tensorclass-eager] 0.2554ms 21.4320μs 46.6593 KOps/s 46.2886 KOps/s $\color{#35bf28}+0.80\%$
test_compile_indexing[slice-pytree-compile] 0.1159ms 71.9091μs 13.9065 KOps/s 13.9306 KOps/s $\color{#d91a1a}-0.17\%$
test_compile_indexing[slice-pytree-eager] 0.2312ms 21.4946μs 46.5234 KOps/s 46.0757 KOps/s $\color{#35bf28}+0.97\%$
test_compile_indexing[int-tensordict-compile] 0.1703ms 88.1447μs 11.3450 KOps/s 11.4765 KOps/s $\color{#d91a1a}-1.15\%$
test_compile_indexing[int-tensordict-eager] 0.2789ms 27.3979μs 36.4991 KOps/s 35.8056 KOps/s $\color{#35bf28}+1.94\%$
test_compile_indexing[int-tensorclass-compile] 0.2065ms 71.6881μs 13.9493 KOps/s 13.9441 KOps/s $\color{#35bf28}+0.04\%$
test_compile_indexing[int-tensorclass-eager] 0.2484ms 21.3466μs 46.8459 KOps/s 46.0321 KOps/s $\color{#35bf28}+1.77\%$
test_compile_indexing[int-pytree-compile] 0.1124ms 69.5097μs 14.3865 KOps/s 13.7781 KOps/s $\color{#35bf28}+4.42\%$
test_compile_indexing[int-pytree-eager] 0.2580ms 21.4565μs 46.6060 KOps/s 47.0863 KOps/s $\color{#d91a1a}-1.02\%$
test_compile_replace[single-eager] 0.1032ms 46.9151μs 21.3151 KOps/s 20.6067 KOps/s $\color{#35bf28}+3.44\%$
test_compile_replace[single-compile] 0.2285ms 0.1327ms 7.5334 KOps/s 7.4556 KOps/s $\color{#35bf28}+1.04\%$
test_compile_replace[multi-eager] 0.7982ms 0.5531ms 1.8079 KOps/s 1.7393 KOps/s $\color{#35bf28}+3.94\%$
test_compile_replace[multi-compile] 0.2944ms 0.1399ms 7.1501 KOps/s 6.9939 KOps/s $\color{#35bf28}+2.23\%$
test_compile_tc_getattr_20[eager] 0.2805ms 0.1651ms 6.0561 KOps/s 6.1706 KOps/s $\color{#d91a1a}-1.86\%$
test_compile_tc_getattr_20[compile] 0.4313ms 0.1464ms 6.8291 KOps/s 6.7379 KOps/s $\color{#35bf28}+1.35\%$
test_compile_clone_shallow[20-eager] 54.9710μs 18.7204μs 53.4175 KOps/s 53.8734 KOps/s $\color{#d91a1a}-0.85\%$
test_compile_clone_shallow[20-compile] 91.2420μs 16.4994μs 60.6084 KOps/s 59.1597 KOps/s $\color{#35bf28}+2.45\%$
test_compile_clone_shallow[40-eager] 87.8610μs 32.7414μs 30.5424 KOps/s 30.1879 KOps/s $\color{#35bf28}+1.17\%$
test_compile_clone_shallow[40-compile] 79.4410μs 18.5089μs 54.0280 KOps/s 55.6358 KOps/s $\color{#d91a1a}-2.89\%$
test_compile_clone_shallow[80-eager] 0.1193ms 61.1474μs 16.3539 KOps/s 16.1322 KOps/s $\color{#35bf28}+1.37\%$
test_compile_clone_shallow[80-compile] 80.1610μs 21.6684μs 46.1502 KOps/s 45.8848 KOps/s $\color{#35bf28}+0.58\%$
test_compile_update_inplace[eager] 0.1143ms 59.1055μs 16.9189 KOps/s 16.9386 KOps/s $\color{#d91a1a}-0.12\%$
test_compile_update_inplace[compile] 0.4480ms 0.1415ms 7.0687 KOps/s 6.9496 KOps/s $\color{#35bf28}+1.71\%$
test_mod_add[eager] 0.3184ms 0.2224ms 4.4959 KOps/s 4.3845 KOps/s $\color{#35bf28}+2.54\%$
test_mod_add[compile] 0.2105ms 0.1295ms 7.7194 KOps/s 7.4972 KOps/s $\color{#35bf28}+2.96\%$
test_mod_add[compile-overhead] 0.3239ms 0.1766ms 5.6639 KOps/s 5.4485 KOps/s $\color{#35bf28}+3.95\%$
test_mod_wrap[eager] 0.6268ms 0.4658ms 2.1471 KOps/s 2.0936 KOps/s $\color{#35bf28}+2.55\%$
test_mod_wrap[compile] 0.4830ms 0.4114ms 2.4305 KOps/s 2.4216 KOps/s $\color{#35bf28}+0.37\%$
test_mod_wrap[compile-overhead] 8.7399ms 4.8066ms 208.0462 Ops/s 210.5514 Ops/s $\color{#d91a1a}-1.19\%$
test_mod_wrap_and_backward[eager] 1.8158ms 1.7043ms 586.7555 Ops/s 591.8626 Ops/s $\color{#d91a1a}-0.86\%$
test_mod_wrap_and_backward[compile] 2.8519ms 1.6755ms 596.8269 Ops/s 655.2446 Ops/s $\textbf{\color{#d91a1a}-8.92\%}$
test_mod_wrap_and_backward[compile-overhead] 1.5083ms 1.0395ms 961.9726 Ops/s 963.7831 Ops/s $\color{#d91a1a}-0.19\%$
test_seq_add[eager] 0.9986ms 0.4979ms 2.0082 KOps/s 2.0161 KOps/s $\color{#d91a1a}-0.39\%$
test_seq_add[compile] 0.6183ms 0.1373ms 7.2837 KOps/s 6.9298 KOps/s $\textbf{\color{#35bf28}+5.11\%}$
test_seq_add[compile-overhead] 0.2988ms 0.1866ms 5.3586 KOps/s 5.2999 KOps/s $\color{#35bf28}+1.11\%$
test_seq_wrap[eager] 1.8247ms 1.3854ms 721.7963 Ops/s 716.5646 Ops/s $\color{#35bf28}+0.73\%$
test_seq_wrap[compile] 0.4989ms 0.4032ms 2.4802 KOps/s 2.4129 KOps/s $\color{#35bf28}+2.79\%$
test_seq_wrap[compile-overhead] 0.4240ms 0.2997ms 3.3364 KOps/s 3.2964 KOps/s $\color{#35bf28}+1.21\%$
test_func_call_runtime[False-eager] 0.9029ms 0.7912ms 1.2639 KOps/s 1.2752 KOps/s $\color{#d91a1a}-0.89\%$
test_func_call_runtime[False-compile] 1.1536ms 0.9577ms 1.0442 KOps/s 1.0363 KOps/s $\color{#35bf28}+0.76\%$
test_func_call_runtime[False-compile-overhead] 0.9393ms 0.4943ms 2.0231 KOps/s 1.9914 KOps/s $\color{#35bf28}+1.59\%$
test_func_call_runtime[True-eager] 1.4956ms 1.0252ms 975.4096 Ops/s 968.7485 Ops/s $\color{#35bf28}+0.69\%$
test_func_call_runtime[True-compile] 1.0847ms 0.9766ms 1.0239 KOps/s 1.0190 KOps/s $\color{#35bf28}+0.49\%$
test_func_call_runtime[True-compile-overhead] 0.6907ms 0.5109ms 1.9573 KOps/s 1.9268 KOps/s $\color{#35bf28}+1.59\%$
test_func_call_cm_runtime[False-eager] 0.9006ms 0.7934ms 1.2604 KOps/s 1.2397 KOps/s $\color{#35bf28}+1.67\%$
test_func_call_cm_runtime[False-compile] 1.1148ms 0.9469ms 1.0561 KOps/s 1.0522 KOps/s $\color{#35bf28}+0.37\%$
test_func_call_cm_runtime[False-compile-overhead] 0.7078ms 0.4778ms 2.0930 KOps/s 2.0735 KOps/s $\color{#35bf28}+0.94\%$
test_func_call_cm_runtime[True-eager] 1.3755ms 1.1730ms 852.5096 Ops/s 858.3494 Ops/s $\color{#d91a1a}-0.68\%$
test_func_call_cm_runtime[True-compile] 1.1317ms 1.0007ms 999.3134 Ops/s 997.7329 Ops/s $\color{#35bf28}+0.16\%$
test_func_call_cm_runtime[True-compile-overhead] 0.6843ms 0.5356ms 1.8672 KOps/s 1.8572 KOps/s $\color{#35bf28}+0.54\%$
test_vmap_func_call_cm_runtime[eager] 2.7766ms 2.3565ms 424.3657 Ops/s 426.7855 Ops/s $\color{#d91a1a}-0.57\%$
test_vmap_func_call_cm_runtime[compile] 1.1537ms 1.0089ms 991.1640 Ops/s 957.7113 Ops/s $\color{#35bf28}+3.49\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.6368ms 0.5240ms 1.9085 KOps/s 1.8979 KOps/s $\color{#35bf28}+0.56\%$
test_distributed 0.8334ms 0.1517ms 6.5913 KOps/s 5.8064 KOps/s $\textbf{\color{#35bf28}+13.52\%}$
test_tdmodule 0.6357ms 0.1904ms 5.2512 KOps/s 5.2300 KOps/s $\color{#35bf28}+0.41\%$
test_tdmodule_dispatch 0.4358ms 0.3644ms 2.7443 KOps/s 2.7445 KOps/s $-0.01\%$
test_tdseq 46.0610μs 26.6318μs 37.5491 KOps/s 38.0325 KOps/s $\color{#d91a1a}-1.27\%$
test_tdseq_dispatch 0.2955ms 0.2115ms 4.7275 KOps/s 4.7207 KOps/s $\color{#35bf28}+0.14\%$
test_instantiation_functorch 2.1799ms 1.9900ms 502.5072 Ops/s 501.9342 Ops/s $\color{#35bf28}+0.11\%$
test_exec_functorch 0.2356ms 0.1801ms 5.5526 KOps/s 5.7039 KOps/s $\color{#d91a1a}-2.65\%$
test_exec_functional_call 0.2241ms 0.1603ms 6.2400 KOps/s 6.2954 KOps/s $\color{#d91a1a}-0.88\%$
test_exec_td_decorator 0.4652ms 0.2331ms 4.2898 KOps/s 4.2614 KOps/s $\color{#35bf28}+0.67\%$
test_vmap_mlp_speed_decorator[True-True] 1.2095ms 1.0093ms 990.8220 Ops/s 1.0006 KOps/s $\color{#d91a1a}-0.98\%$
test_vmap_mlp_speed_decorator[True-False] 1.1991ms 1.0060ms 994.0532 Ops/s 999.3949 Ops/s $\color{#d91a1a}-0.53\%$
test_vmap_mlp_speed_decorator[False-True] 0.9268ms 0.7139ms 1.4007 KOps/s 1.4169 KOps/s $\color{#d91a1a}-1.14\%$
test_vmap_mlp_speed_decorator[False-False] 0.8830ms 0.7106ms 1.4073 KOps/s 1.4107 KOps/s $\color{#d91a1a}-0.24\%$
test_vmap_transformer_speed_decorator[True-True] 21.5790ms 20.9053ms 47.8347 Ops/s 48.4150 Ops/s $\color{#d91a1a}-1.20\%$
test_vmap_transformer_speed_decorator[True-False] 21.4500ms 20.8691ms 47.9177 Ops/s 48.3578 Ops/s $\color{#d91a1a}-0.91\%$
test_vmap_transformer_speed_decorator[False-True] 21.6794ms 20.5460ms 48.6713 Ops/s 49.1981 Ops/s $\color{#d91a1a}-1.07\%$
test_vmap_transformer_speed_decorator[False-False] 21.1855ms 20.4688ms 48.8548 Ops/s 49.5425 Ops/s $\color{#d91a1a}-1.39\%$
test_to_module_speed[True] 2.1063ms 1.4206ms 703.9374 Ops/s 699.8330 Ops/s $\color{#35bf28}+0.59\%$
test_to_module_speed[False] 1.6159ms 1.3909ms 718.9731 Ops/s 708.2378 Ops/s $\color{#35bf28}+1.52\%$
test_tc_init 75.0210μs 40.8996μs 24.4501 KOps/s 23.8655 KOps/s $\color{#35bf28}+2.45\%$
test_tc_init_tensor_only 43.0200μs 8.5293μs 117.2431 KOps/s 117.9727 KOps/s $\color{#d91a1a}-0.62\%$
test_tc_init_nested 0.1175ms 81.6741μs 12.2438 KOps/s 12.0407 KOps/s $\color{#35bf28}+1.69\%$
test_tc_init_many_fields 57.4210μs 14.0499μs 71.1749 KOps/s 70.7806 KOps/s $\color{#35bf28}+0.56\%$
test_tc_first_layer_tensor 29.5900μs 1.6999μs 588.2709 KOps/s 582.5242 KOps/s $\color{#35bf28}+0.99\%$
test_tc_first_layer_tensor_only 2.6161μs 0.3866μs 2.5867 MOps/s 2.5606 MOps/s $\color{#35bf28}+1.02\%$
test_tc_first_layer_tensor_set 21.6610μs 3.7299μs 268.1007 KOps/s 271.9010 KOps/s $\color{#d91a1a}-1.40\%$
test_tc_first_layer_tensor_only_set 25.8800μs 3.1508μs 317.3762 KOps/s 321.6758 KOps/s $\color{#d91a1a}-1.34\%$
test_tc_first_layer_nontensor 41.3600μs 5.8599μs 170.6515 KOps/s 166.0254 KOps/s $\color{#35bf28}+2.79\%$
test_tc_second_layer_tensor 38.9210μs 4.1397μs 241.5638 KOps/s 240.2813 KOps/s $\color{#35bf28}+0.53\%$
test_tc_second_layer_nontensor 33.8510μs 8.2001μs 121.9493 KOps/s 119.3822 KOps/s $\color{#35bf28}+2.15\%$
test_unbind 0.2805s 12.8209ms 77.9976 Ops/s 62.8578 Ops/s $\textbf{\color{#35bf28}+24.09\%}$
test_full_like 4.7976ms 4.3593ms 229.3924 Ops/s 231.3566 Ops/s $\color{#d91a1a}-0.85\%$
test_zeros_like 4.6555ms 4.2770ms 233.8108 Ops/s 232.1435 Ops/s $\color{#35bf28}+0.72\%$
test_ones_like 4.4268ms 4.2815ms 233.5635 Ops/s 233.5798 Ops/s $-0.01\%$
test_clone 6.8331ms 6.3966ms 156.3327 Ops/s 157.0859 Ops/s $\color{#d91a1a}-0.48\%$
test_squeeze 74.4010μs 14.3904μs 69.4906 KOps/s 71.5791 KOps/s $\color{#d91a1a}-2.92\%$
test_unsqueeze 0.1678ms 0.1104ms 9.0578 KOps/s 8.8260 KOps/s $\color{#35bf28}+2.63\%$
test_split 0.2586ms 0.1769ms 5.6527 KOps/s 5.6343 KOps/s $\color{#35bf28}+0.33\%$
test_permute 0.2821ms 0.2112ms 4.7344 KOps/s 4.9190 KOps/s $\color{#d91a1a}-3.75\%$
test_stack 51.3913ms 51.1141ms 19.5641 Ops/s 19.5307 Ops/s $\color{#35bf28}+0.17\%$
test_cat 51.3228ms 50.9824ms 19.6146 Ops/s 19.5052 Ops/s $\color{#35bf28}+0.56\%$
test_sequential_tensordict 0.6223ms 0.5624ms 1.7780 KOps/s 1.7925 KOps/s $\color{#d91a1a}-0.81\%$
test_sequential_graph_module 0.1767ms 0.1229ms 8.1354 KOps/s 8.5375 KOps/s $\color{#d91a1a}-4.71\%$
test_nested_tensordict 0.8902ms 0.8145ms 1.2278 KOps/s 1.2582 KOps/s $\color{#d91a1a}-2.42\%$
test_nested_graph_module 0.1993ms 0.1331ms 7.5108 KOps/s 7.9242 KOps/s $\textbf{\color{#d91a1a}-5.22\%}$

@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}20$. Worsened: $\large\color{#d91a1a}9$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_collector_preallocated_write[stack_out-cpu-32-2] 0.2170ms 0.1665ms 6.0054 KOps/s 5.7915 KOps/s $\color{#35bf28}+3.69\%$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.3509ms 0.2958ms 3.3801 KOps/s 3.3702 KOps/s $\color{#35bf28}+0.30\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8441ms 0.7744ms 1.2914 KOps/s 1.2804 KOps/s $\color{#35bf28}+0.85\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.4345ms 0.3333ms 3.0005 KOps/s 2.9900 KOps/s $\color{#35bf28}+0.35\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.2989ms 1.1725ms 852.9121 Ops/s 928.8797 Ops/s $\textbf{\color{#d91a1a}-8.18\%}$
test_collector_preallocated_write[stack_out-cpu-4096-32] 7.4178ms 7.2168ms 138.5647 Ops/s 135.5938 Ops/s $\color{#35bf28}+2.19\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.2966ms 0.2477ms 4.0371 KOps/s 3.8909 KOps/s $\color{#35bf28}+3.76\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.4514ms 0.3709ms 2.6959 KOps/s 2.5899 KOps/s $\color{#35bf28}+4.09\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.9274ms 0.8226ms 1.2157 KOps/s 1.2090 KOps/s $\color{#35bf28}+0.56\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.3316ms 0.2452ms 4.0789 KOps/s 3.7530 KOps/s $\textbf{\color{#35bf28}+8.68\%}$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.4962ms 0.3695ms 2.7061 KOps/s 2.6648 KOps/s $\color{#35bf28}+1.55\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 0.9733ms 0.8865ms 1.1281 KOps/s 1.1333 KOps/s $\color{#d91a1a}-0.46\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.3142ms 0.2135ms 4.6838 KOps/s 4.5621 KOps/s $\color{#35bf28}+2.67\%$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 1.6388ms 0.8327ms 1.2009 KOps/s 1.1656 KOps/s $\color{#35bf28}+3.03\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.4317ms 3.3124ms 301.8997 Ops/s 297.8015 Ops/s $\color{#35bf28}+1.38\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.4662ms 0.3696ms 2.7057 KOps/s 2.6876 KOps/s $\color{#35bf28}+0.67\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 1.9658ms 1.7692ms 565.2284 Ops/s 603.8624 Ops/s $\textbf{\color{#d91a1a}-6.40\%}$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 10.9174ms 10.2999ms 97.0885 Ops/s 97.1658 Ops/s $\color{#d91a1a}-0.08\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.4559ms 0.3954ms 2.5289 KOps/s 2.5194 KOps/s $\color{#35bf28}+0.38\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.5406ms 1.4332ms 697.7599 Ops/s 673.8941 Ops/s $\color{#35bf28}+3.54\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 5.8786ms 5.6652ms 176.5166 Ops/s 171.5581 Ops/s $\color{#35bf28}+2.89\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.4686ms 0.3872ms 2.5825 KOps/s 2.5184 KOps/s $\color{#35bf28}+2.54\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.5334ms 1.4514ms 688.9914 Ops/s 668.5043 Ops/s $\color{#35bf28}+3.06\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 5.9567ms 5.7207ms 174.8025 Ops/s 170.9166 Ops/s $\color{#35bf28}+2.27\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.2387ms 0.1741ms 5.7454 KOps/s 5.7895 KOps/s $\color{#d91a1a}-0.76\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 0.7609ms 0.6695ms 1.4936 KOps/s 1.4701 KOps/s $\color{#35bf28}+1.59\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 2.7578ms 2.6401ms 378.7696 Ops/s 370.2263 Ops/s $\color{#35bf28}+2.31\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.4981ms 0.3146ms 3.1789 KOps/s 3.1151 KOps/s $\color{#35bf28}+2.05\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 1.6722ms 1.5758ms 634.5950 Ops/s 690.4487 Ops/s $\textbf{\color{#d91a1a}-8.09\%}$
test_collector_preallocated_write[copy_at-cpu-4096-32] 9.9271ms 9.5040ms 105.2188 Ops/s 106.1377 Ops/s $\color{#d91a1a}-0.87\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.3949ms 0.3232ms 3.0941 KOps/s 3.0683 KOps/s $\color{#35bf28}+0.84\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.3005ms 1.2056ms 829.4628 Ops/s 821.3613 Ops/s $\color{#35bf28}+0.99\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 4.8478ms 4.7067ms 212.4642 Ops/s 209.8803 Ops/s $\color{#35bf28}+1.23\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.3865ms 0.3256ms 3.0713 KOps/s 3.0425 KOps/s $\color{#35bf28}+0.94\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.3137ms 1.2099ms 826.5023 Ops/s 817.1197 Ops/s $\color{#35bf28}+1.15\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 4.8324ms 4.7257ms 211.6073 Ops/s 207.6117 Ops/s $\color{#35bf28}+1.92\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1491ms 0.1195ms 8.3688 KOps/s 8.2552 KOps/s $\color{#35bf28}+1.38\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5578ms 0.4701ms 2.1271 KOps/s 2.1190 KOps/s $\color{#35bf28}+0.39\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 1.9606ms 1.8753ms 533.2534 Ops/s 532.2036 Ops/s $\color{#35bf28}+0.20\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.3402ms 0.2484ms 4.0250 KOps/s 3.9091 KOps/s $\color{#35bf28}+2.97\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.4082ms 1.2289ms 813.7079 Ops/s 814.2905 Ops/s $\color{#d91a1a}-0.07\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 8.6710ms 8.5357ms 117.1546 Ops/s 117.7944 Ops/s $\color{#d91a1a}-0.54\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.3473ms 0.2743ms 3.6457 KOps/s 3.6091 KOps/s $\color{#35bf28}+1.01\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.0987ms 1.0237ms 976.8499 Ops/s 962.2875 Ops/s $\color{#35bf28}+1.51\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.1650ms 3.9757ms 251.5273 Ops/s 246.9315 Ops/s $\color{#35bf28}+1.86\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.3479ms 0.2744ms 3.6440 KOps/s 3.6101 KOps/s $\color{#35bf28}+0.94\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.0858ms 1.0197ms 980.6503 Ops/s 960.9556 Ops/s $\color{#35bf28}+2.05\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.1123ms 3.9920ms 250.4984 Ops/s 245.1846 Ops/s $\color{#35bf28}+2.17\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 16.6852ms 16.4358ms 60.8429 Ops/s 61.4588 Ops/s $\color{#d91a1a}-1.00\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 33.4230ms 33.1375ms 30.1773 Ops/s 30.3518 Ops/s $\color{#d91a1a}-0.58\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 85.4043ms 84.1788ms 11.8795 Ops/s 12.1226 Ops/s $\color{#d91a1a}-2.01\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2018s 0.2012s 4.9708 Ops/s 5.0904 Ops/s $\color{#d91a1a}-2.35\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 14.0636ms 13.7846ms 72.5445 Ops/s 72.6506 Ops/s $\color{#d91a1a}-0.15\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 27.2326ms 26.8926ms 37.1849 Ops/s 37.3280 Ops/s $\color{#d91a1a}-0.38\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 14.0733ms 13.9173ms 71.8528 Ops/s 71.7203 Ops/s $\color{#35bf28}+0.18\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 28.6131ms 27.4295ms 36.4572 Ops/s 36.5516 Ops/s $\color{#d91a1a}-0.26\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 52.6437ms 52.2907ms 19.1239 Ops/s 18.9019 Ops/s $\color{#35bf28}+1.17\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1069s 0.1059s 9.4453 Ops/s 9.4079 Ops/s $\color{#35bf28}+0.40\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1437s 0.1433s 6.9781 Ops/s 6.9063 Ops/s $\color{#35bf28}+1.04\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3281s 0.3272s 3.0558 Ops/s 3.1440 Ops/s $\color{#d91a1a}-2.80\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 75.2774ms 74.2939ms 13.4601 Ops/s 13.3453 Ops/s $\color{#35bf28}+0.86\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1499s 0.1485s 6.7318 Ops/s 6.6252 Ops/s $\color{#35bf28}+1.61\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 75.5570ms 75.1226ms 13.3116 Ops/s 13.2101 Ops/s $\color{#35bf28}+0.77\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1511s 0.1501s 6.6641 Ops/s 6.6151 Ops/s $\color{#35bf28}+0.74\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.7361ms 30.4395ms 32.8520 Ops/s 32.8379 Ops/s $\color{#35bf28}+0.04\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 62.5219ms 61.9886ms 16.1320 Ops/s 16.1837 Ops/s $\color{#d91a1a}-0.32\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1203s 0.1194s 8.3767 Ops/s 8.5896 Ops/s $\color{#d91a1a}-2.48\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2732s 0.2724s 3.6717 Ops/s 3.8370 Ops/s $\color{#d91a1a}-4.31\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 43.7054ms 43.2091ms 23.1433 Ops/s 22.9007 Ops/s $\color{#35bf28}+1.06\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 88.5910ms 86.8770ms 11.5105 Ops/s 11.5318 Ops/s $\color{#d91a1a}-0.18\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 44.0013ms 43.6126ms 22.9291 Ops/s 22.8316 Ops/s $\color{#35bf28}+0.43\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 87.8078ms 87.3344ms 11.4502 Ops/s 11.3876 Ops/s $\color{#35bf28}+0.55\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 20.6942ms 20.5013ms 48.7774 Ops/s 48.5342 Ops/s $\color{#35bf28}+0.50\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 42.0847ms 41.6922ms 23.9853 Ops/s 23.9854 Ops/s $-0.00\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1063s 0.1057s 9.4635 Ops/s 9.5162 Ops/s $\color{#d91a1a}-0.55\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2490s 0.2487s 4.0206 Ops/s 4.2124 Ops/s $\color{#d91a1a}-4.55\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 39.7020ms 38.6512ms 25.8724 Ops/s 26.8064 Ops/s $\color{#d91a1a}-3.48\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 77.7348ms 77.0083ms 12.9856 Ops/s 13.3464 Ops/s $\color{#d91a1a}-2.70\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 39.3502ms 38.6879ms 25.8479 Ops/s 26.9459 Ops/s $\color{#d91a1a}-4.07\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 78.1886ms 77.1922ms 12.9547 Ops/s 13.4512 Ops/s $\color{#d91a1a}-3.69\%$
test_plain_set_nested 38.2510μs 15.4661μs 64.6576 KOps/s 64.8198 KOps/s $\color{#d91a1a}-0.25\%$
test_plain_set_stack_nested 43.3110μs 15.2607μs 65.5276 KOps/s 65.6313 KOps/s $\color{#d91a1a}-0.16\%$
test_plain_set_nested_inplace 41.7910μs 16.8314μs 59.4128 KOps/s 58.8349 KOps/s $\color{#35bf28}+0.98\%$
test_plain_set_stack_nested_inplace 64.3020μs 16.5950μs 60.2591 KOps/s 60.0155 KOps/s $\color{#35bf28}+0.41\%$
test_items 35.5010μs 6.0748μs 164.6141 KOps/s 166.8053 KOps/s $\color{#d91a1a}-1.31\%$
test_items_nested 0.6110ms 0.4643ms 2.1539 KOps/s 2.1847 KOps/s $\color{#d91a1a}-1.41\%$
test_items_nested_locked 0.5254ms 0.4662ms 2.1451 KOps/s 2.1644 KOps/s $\color{#d91a1a}-0.89\%$
test_items_nested_leaf 0.1354ms 97.6440μs 10.2413 KOps/s 10.2733 KOps/s $\color{#d91a1a}-0.31\%$
test_items_stack_nested 0.5201ms 0.4633ms 2.1585 KOps/s 2.1750 KOps/s $\color{#d91a1a}-0.76\%$
test_items_stack_nested_leaf 0.1343ms 97.5741μs 10.2486 KOps/s 9.9993 KOps/s $\color{#35bf28}+2.49\%$
test_items_stack_nested_locked 0.5350ms 0.4647ms 2.1518 KOps/s 2.1524 KOps/s $\color{#d91a1a}-0.03\%$
test_keys 30.9610μs 4.2393μs 235.8868 KOps/s 235.6866 KOps/s $\color{#35bf28}+0.08\%$
test_keys_nested 0.1822ms 0.1311ms 7.6277 KOps/s 7.5894 KOps/s $\color{#35bf28}+0.50\%$
test_keys_nested_locked 0.7114ms 0.1400ms 7.1446 KOps/s 7.1640 KOps/s $\color{#d91a1a}-0.27\%$
test_keys_nested_leaf 0.1766ms 0.1210ms 8.2662 KOps/s 8.1832 KOps/s $\color{#35bf28}+1.01\%$
test_keys_stack_nested 0.1834ms 0.1306ms 7.6582 KOps/s 7.6604 KOps/s $\color{#d91a1a}-0.03\%$
test_keys_stack_nested_leaf 0.1653ms 0.1197ms 8.3535 KOps/s 8.2478 KOps/s $\color{#35bf28}+1.28\%$
test_keys_stack_nested_locked 0.2078ms 0.1400ms 7.1410 KOps/s 7.2092 KOps/s $\color{#d91a1a}-0.95\%$
test_values 10.3042μs 1.0273μs 973.4562 KOps/s 986.6963 KOps/s $\color{#d91a1a}-1.34\%$
test_values_nested 86.0520μs 53.0669μs 18.8441 KOps/s 18.8648 KOps/s $\color{#d91a1a}-0.11\%$
test_values_nested_locked 87.6120μs 56.7143μs 17.6322 KOps/s 17.5749 KOps/s $\color{#35bf28}+0.33\%$
test_values_nested_leaf 91.1530μs 61.0336μs 16.3844 KOps/s 16.5096 KOps/s $\color{#d91a1a}-0.76\%$
test_values_stack_nested 91.2630μs 52.7053μs 18.9734 KOps/s 18.7818 KOps/s $\color{#35bf28}+1.02\%$
test_values_stack_nested_leaf 98.2220μs 60.6420μs 16.4902 KOps/s 16.5837 KOps/s $\color{#d91a1a}-0.56\%$
test_values_stack_nested_locked 89.9920μs 56.7402μs 17.6242 KOps/s 17.6379 KOps/s $\color{#d91a1a}-0.08\%$
test_membership 4.4018μs 0.8478μs 1.1796 MOps/s 1.1743 MOps/s $\color{#35bf28}+0.45\%$
test_membership_nested 35.3610μs 2.9835μs 335.1740 KOps/s 333.4694 KOps/s $\color{#35bf28}+0.51\%$
test_membership_nested_leaf 26.0000μs 2.9938μs 334.0190 KOps/s 334.3967 KOps/s $\color{#d91a1a}-0.11\%$
test_membership_stacked_nested 41.5210μs 2.8793μs 347.3006 KOps/s 345.3641 KOps/s $\color{#35bf28}+0.56\%$
test_membership_stacked_nested_leaf 34.2410μs 2.8876μs 346.3030 KOps/s 346.9686 KOps/s $\color{#d91a1a}-0.19\%$
test_membership_nested_last 50.3810μs 4.4538μs 224.5254 KOps/s 222.5394 KOps/s $\color{#35bf28}+0.89\%$
test_membership_nested_leaf_last 35.1910μs 4.4225μs 226.1172 KOps/s 221.7147 KOps/s $\color{#35bf28}+1.99\%$
test_membership_stacked_nested_last 39.7920μs 4.3432μs 230.2446 KOps/s 228.3040 KOps/s $\color{#35bf28}+0.85\%$
test_membership_stacked_nested_leaf_last 29.0210μs 4.3641μs 229.1414 KOps/s 230.5384 KOps/s $\color{#d91a1a}-0.61\%$
test_nested_getleaf 49.1610μs 22.0760μs 45.2980 KOps/s 45.4463 KOps/s $\color{#d91a1a}-0.33\%$
test_nested_get 44.2910μs 20.8342μs 47.9980 KOps/s 48.1582 KOps/s $\color{#d91a1a}-0.33\%$
test_stacked_getleaf 52.9010μs 21.7902μs 45.8923 KOps/s 45.2827 KOps/s $\color{#35bf28}+1.35\%$
test_stacked_get 47.7010μs 20.6138μs 48.5112 KOps/s 47.7931 KOps/s $\color{#35bf28}+1.50\%$
test_nested_getitemleaf 49.0510μs 22.4450μs 44.5534 KOps/s 44.6599 KOps/s $\color{#d91a1a}-0.24\%$
test_nested_getitem 45.7810μs 21.2399μs 47.0812 KOps/s 47.1459 KOps/s $\color{#d91a1a}-0.14\%$
test_stacked_getitemleaf 49.1610μs 22.4338μs 44.5756 KOps/s 44.6244 KOps/s $\color{#d91a1a}-0.11\%$
test_stacked_getitem 45.8010μs 21.1949μs 47.1812 KOps/s 46.1651 KOps/s $\color{#35bf28}+2.20\%$
test_lock_nested 0.6640ms 0.5008ms 1.9969 KOps/s 2.0862 KOps/s $\color{#d91a1a}-4.28\%$
test_lock_stack_nested 0.5822ms 0.5080ms 1.9683 KOps/s 2.0928 KOps/s $\textbf{\color{#d91a1a}-5.95\%}$
test_unlock_nested 0.4633ms 0.3977ms 2.5146 KOps/s 2.5961 KOps/s $\color{#d91a1a}-3.14\%$
test_unlock_stack_nested 0.5154ms 0.4032ms 2.4799 KOps/s 2.5771 KOps/s $\color{#d91a1a}-3.77\%$
test_flatten_speed 0.2291ms 0.1225ms 8.1622 KOps/s 8.1729 KOps/s $\color{#d91a1a}-0.13\%$
test_unflatten_speed 0.6506ms 0.5788ms 1.7276 KOps/s 1.7473 KOps/s $\color{#d91a1a}-1.13\%$
test_common_ops 0.1001s 1.5079ms 663.1896 Ops/s 324.7901 Ops/s $\textbf{\color{#35bf28}+104.19\%}$
test_creation 0.1336ms 3.1611μs 316.3425 KOps/s 322.6085 KOps/s $\color{#d91a1a}-1.94\%$
test_creation_empty 31.9210μs 7.0099μs 142.6558 KOps/s 141.7291 KOps/s $\color{#35bf28}+0.65\%$
test_creation_nested_1 34.3810μs 11.5743μs 86.3983 KOps/s 86.5718 KOps/s $\color{#d91a1a}-0.20\%$
test_creation_nested_2 43.4010μs 13.8523μs 72.1902 KOps/s 73.7015 KOps/s $\color{#d91a1a}-2.05\%$
test_creation_many_keys[10] 52.4810μs 21.1449μs 47.2927 KOps/s 47.3959 KOps/s $\color{#d91a1a}-0.22\%$
test_creation_many_keys[50] 0.1311ms 91.0086μs 10.9880 KOps/s 10.9904 KOps/s $\color{#d91a1a}-0.02\%$
test_creation_many_keys[100] 0.2411ms 0.1768ms 5.6569 KOps/s 5.5629 KOps/s $\color{#35bf28}+1.69\%$
test_creation_nested_many_keys[10] 73.7420μs 45.1471μs 22.1498 KOps/s 21.8066 KOps/s $\color{#35bf28}+1.57\%$
test_creation_nested_many_keys[50] 0.2298ms 0.1854ms 5.3941 KOps/s 5.3502 KOps/s $\color{#35bf28}+0.82\%$
test_clone 31.9510μs 13.4682μs 74.2488 KOps/s 74.6228 KOps/s $\color{#d91a1a}-0.50\%$
test_getitem[int] 1.5187ms 15.4429μs 64.7547 KOps/s 65.3555 KOps/s $\color{#d91a1a}-0.92\%$
test_getitem[slice_int] 0.1461ms 25.6866μs 38.9308 KOps/s 39.7689 KOps/s $\color{#d91a1a}-2.11\%$
test_getitem[range] 0.1900ms 65.0851μs 15.3645 KOps/s 15.3232 KOps/s $\color{#35bf28}+0.27\%$
test_getitem[tuple] 0.1429ms 24.6966μs 40.4914 KOps/s 40.6189 KOps/s $\color{#d91a1a}-0.31\%$
test_getitem[list] 0.1872ms 59.4520μs 16.8203 KOps/s 16.6785 KOps/s $\color{#35bf28}+0.85\%$
test_setitem_dim[int] 56.1010μs 25.1016μs 39.8381 KOps/s 38.6037 KOps/s $\color{#35bf28}+3.20\%$
test_setitem_dim[slice_int] 64.6620μs 42.1145μs 23.7448 KOps/s 23.0422 KOps/s $\color{#35bf28}+3.05\%$
test_setitem_dim[range] 0.1277ms 95.8080μs 10.4375 KOps/s 10.3633 KOps/s $\color{#35bf28}+0.72\%$
test_setitem_dim[tuple] 59.3610μs 38.9964μs 25.6434 KOps/s 25.4874 KOps/s $\color{#35bf28}+0.61\%$
test_setitem 45.4510μs 17.6619μs 56.6191 KOps/s 55.8974 KOps/s $\color{#35bf28}+1.29\%$
test_set 41.2710μs 17.0138μs 58.7758 KOps/s 58.9347 KOps/s $\color{#d91a1a}-0.27\%$
test_set_shared 0.6271ms 0.2103ms 4.7550 KOps/s 4.8696 KOps/s $\color{#d91a1a}-2.35\%$
test_update 0.2827ms 21.9874μs 45.4805 KOps/s 45.7182 KOps/s $\color{#d91a1a}-0.52\%$
test_update_nested 58.4420μs 33.3182μs 30.0137 KOps/s 29.5860 KOps/s $\color{#35bf28}+1.45\%$
test_update__nested 0.4638ms 33.6319μs 29.7337 KOps/s 30.7085 KOps/s $\color{#d91a1a}-3.17\%$
test_set_nested 56.2310μs 18.8196μs 53.1360 KOps/s 53.4096 KOps/s $\color{#d91a1a}-0.51\%$
test_set_nested_new 60.9210μs 23.9655μs 41.7267 KOps/s 42.1627 KOps/s $\color{#d91a1a}-1.03\%$
test_select 63.0110μs 40.4599μs 24.7158 KOps/s 24.8232 KOps/s $\color{#d91a1a}-0.43\%$
test_select_nested 0.1128ms 74.7284μs 13.3818 KOps/s 13.4584 KOps/s $\color{#d91a1a}-0.57\%$
test_exclude_nested 0.1255ms 92.0823μs 10.8599 KOps/s 10.8055 KOps/s $\color{#35bf28}+0.50\%$
test_empty[True] 0.4597ms 0.4019ms 2.4882 KOps/s 2.5005 KOps/s $\color{#d91a1a}-0.49\%$
test_empty[False] 10.1827μs 1.4036μs 712.4692 KOps/s 721.8784 KOps/s $\color{#d91a1a}-1.30\%$
test_to 0.1059ms 69.7380μs 14.3394 KOps/s 14.3236 KOps/s $\color{#35bf28}+0.11\%$
test_to_nonblocking 0.1086ms 69.0213μs 14.4883 KOps/s 14.4743 KOps/s $\color{#35bf28}+0.10\%$
test_unbind_speed 0.4027ms 0.3390ms 2.9496 KOps/s 2.9972 KOps/s $\color{#d91a1a}-1.59\%$
test_unbind_speed_stack0 0.3912ms 0.3362ms 2.9746 KOps/s 3.0269 KOps/s $\color{#d91a1a}-1.73\%$
test_unbind_speed_stack1 0.1154s 0.8662ms 1.1545 KOps/s 1.1651 KOps/s $\color{#d91a1a}-0.91\%$
test_split 1.2621ms 1.1500ms 869.5991 Ops/s 710.0203 Ops/s $\textbf{\color{#35bf28}+22.48\%}$
test_chunk 0.1158s 1.2358ms 809.1671 Ops/s 920.4563 Ops/s $\textbf{\color{#d91a1a}-12.09\%}$
test_to_cpu_blocking 20.1131ms 19.9313ms 50.1724 Ops/s 48.7143 Ops/s $\color{#35bf28}+2.99\%$
test_to_cpu_global_sync 12.1139ms 11.8650ms 84.2818 Ops/s 84.9100 Ops/s $\color{#d91a1a}-0.74\%$
test_to_cpu_event_sync 13.1805ms 12.8205ms 77.9999 Ops/s 77.9164 Ops/s $\color{#35bf28}+0.11\%$
test_to_cpu_default 13.1360ms 12.8233ms 77.9827 Ops/s 69.3721 Ops/s $\textbf{\color{#35bf28}+12.41\%}$
test_consolidate[False-None] 0.1451s 4.8060ms 208.0752 Ops/s 239.3551 Ops/s $\textbf{\color{#d91a1a}-13.07\%}$
test_consolidate[default-None] 3.5332ms 2.1686ms 461.1326 Ops/s 456.7339 Ops/s $\color{#35bf28}+0.96\%$
test_consolidate[reduce-overhead-None] 3.5282ms 2.0831ms 480.0582 Ops/s 469.5448 Ops/s $\color{#35bf28}+2.24\%$
test_consolidate_njt[False-None] 0.2109s 10.2129ms 97.9151 Ops/s 114.8955 Ops/s $\textbf{\color{#d91a1a}-14.78\%}$
test_to[False-False-None] 2.2944ms 2.1891ms 456.8006 Ops/s 454.4380 Ops/s $\color{#35bf28}+0.52\%$
test_to[True-False-None] 2.1335ms 1.9536ms 511.8733 Ops/s 530.7435 Ops/s $\color{#d91a1a}-3.56\%$
test_to[within-False-None] 6.5920ms 6.2495ms 160.0118 Ops/s 162.2715 Ops/s $\color{#d91a1a}-1.39\%$
test_to[True-default-None] 11.0804ms 10.7591ms 92.9446 Ops/s 92.0623 Ops/s $\color{#35bf28}+0.96\%$
test_to_njt[False-False-None] 8.8329ms 8.6702ms 115.3380 Ops/s 113.8424 Ops/s $\color{#35bf28}+1.31\%$
test_to_njt[True-False-None] 7.1101ms 6.9271ms 144.3606 Ops/s 141.0944 Ops/s $\color{#35bf28}+2.31\%$
test_to_njt[within-False-None] 16.0356ms 15.6288ms 63.9845 Ops/s 62.9345 Ops/s $\color{#35bf28}+1.67\%$
test_creation[device0] 0.2961ms 0.1177ms 8.4984 KOps/s 8.1534 KOps/s $\color{#35bf28}+4.23\%$
test_creation_from_tensor 0.4584ms 0.1144ms 8.7447 KOps/s 8.3734 KOps/s $\color{#35bf28}+4.43\%$
test_add_one[memmap_tensor0] 0.3487ms 6.8980μs 144.9698 KOps/s 143.7479 KOps/s $\color{#35bf28}+0.85\%$
test_contiguous[memmap_tensor0] 13.3910μs 0.6596μs 1.5161 MOps/s 2.1500 MOps/s $\textbf{\color{#d91a1a}-29.48\%}$
test_stack[memmap_tensor0] 53.4410μs 4.2608μs 234.7004 KOps/s 228.4292 KOps/s $\color{#35bf28}+2.75\%$
test_memmaptd_index 1.2807ms 0.2774ms 3.6048 KOps/s 3.6210 KOps/s $\color{#d91a1a}-0.45\%$
test_memmaptd_index_astensor 0.5390ms 0.3817ms 2.6197 KOps/s 2.6413 KOps/s $\color{#d91a1a}-0.82\%$
test_memmaptd_index_op 0.7931ms 0.6374ms 1.5690 KOps/s 1.5413 KOps/s $\color{#35bf28}+1.80\%$
test_serialize_model 0.1347s 0.1334s 7.4952 Ops/s 7.4150 Ops/s $\color{#35bf28}+1.08\%$
test_serialize_model_pickle 1.3494s 1.2102s 0.8263 Ops/s 0.8265 Ops/s $\color{#d91a1a}-0.02\%$
test_serialize_weights 0.1356s 0.1335s 7.4925 Ops/s 7.4749 Ops/s $\color{#35bf28}+0.23\%$
test_serialize_weights_returnearly 0.4698s 87.5955ms 11.4161 Ops/s 5.8495 Ops/s $\textbf{\color{#35bf28}+95.16\%}$
test_serialize_weights_pickle 1.3790s 1.2146s 0.8233 Ops/s 0.8217 Ops/s $\color{#35bf28}+0.19\%$
test_reshape_pytree 0.2093ms 33.0611μs 30.2471 KOps/s 30.3021 KOps/s $\color{#d91a1a}-0.18\%$
test_reshape_td 84.5020μs 45.0458μs 22.1996 KOps/s 21.9913 KOps/s $\color{#35bf28}+0.95\%$
test_view_pytree 0.2140ms 32.5519μs 30.7202 KOps/s 30.7660 KOps/s $\color{#d91a1a}-0.15\%$
test_view_td 89.3130μs 52.8659μs 18.9158 KOps/s 18.2740 KOps/s $\color{#35bf28}+3.51\%$
test_unbind_pytree 0.2320ms 36.8373μs 27.1464 KOps/s 26.8535 KOps/s $\color{#35bf28}+1.09\%$
test_unbind_td 95.9020μs 49.8393μs 20.0645 KOps/s 19.6216 KOps/s $\color{#35bf28}+2.26\%$
test_split_pytree 0.1920ms 42.9008μs 23.3096 KOps/s 23.1095 KOps/s $\color{#35bf28}+0.87\%$
test_split_td 0.1145ms 64.3784μs 15.5332 KOps/s 15.2383 KOps/s $\color{#35bf28}+1.93\%$
test_add_pytree 0.2292ms 43.2635μs 23.1142 KOps/s 23.2567 KOps/s $\color{#d91a1a}-0.61\%$
test_add_td 0.1063ms 57.9719μs 17.2497 KOps/s 17.7395 KOps/s $\color{#d91a1a}-2.76\%$
test_compile_add_one_nested[tensordict-compile] 0.3418ms 0.2192ms 4.5611 KOps/s 4.1087 KOps/s $\textbf{\color{#35bf28}+11.01\%}$
test_compile_add_one_nested[tensordict-eager] 0.3352ms 0.2054ms 4.8689 KOps/s 5.0032 KOps/s $\color{#d91a1a}-2.68\%$
test_compile_add_one_nested[pytree-compile] 0.3433ms 0.1359ms 7.3562 KOps/s 7.0131 KOps/s $\color{#35bf28}+4.89\%$
test_compile_add_one_nested[pytree-eager] 0.4364ms 0.1784ms 5.6062 KOps/s 5.5787 KOps/s $\color{#35bf28}+0.49\%$
test_compile_copy_nested[tensordict-compile] 74.3620μs 15.2102μs 65.7452 KOps/s 66.1790 KOps/s $\color{#d91a1a}-0.66\%$
test_compile_copy_nested[tensordict-eager] 0.3067ms 53.5069μs 18.6892 KOps/s 18.2510 KOps/s $\color{#35bf28}+2.40\%$
test_compile_copy_nested[pytree-compile] 0.1266ms 15.7640μs 63.4358 KOps/s 62.5060 KOps/s $\color{#35bf28}+1.49\%$
test_compile_copy_nested[pytree-eager] 0.3826ms 69.0443μs 14.4835 KOps/s 14.6746 KOps/s $\color{#d91a1a}-1.30\%$
test_compile_add_one_flat[tensordict-compile] 0.2949ms 0.2497ms 4.0056 KOps/s 3.8337 KOps/s $\color{#35bf28}+4.49\%$
test_compile_add_one_flat[tensordict-eager] 0.3436ms 0.2831ms 3.5321 KOps/s 3.5501 KOps/s $\color{#d91a1a}-0.51\%$
test_compile_add_one_flat[tensorclass-compile] 0.3897ms 0.1850ms 5.4065 KOps/s 5.2146 KOps/s $\color{#35bf28}+3.68\%$
test_compile_add_one_flat[tensorclass-eager] 0.1192ms 74.5587μs 13.4123 KOps/s 13.3185 KOps/s $\color{#35bf28}+0.70\%$
test_compile_add_one_flat[pytree-compile] 0.2907ms 0.1936ms 5.1663 KOps/s 4.8638 KOps/s $\textbf{\color{#35bf28}+6.22\%}$
test_compile_add_one_flat[pytree-eager] 0.7809ms 0.5132ms 1.9485 KOps/s 1.8972 KOps/s $\color{#35bf28}+2.70\%$
test_compile_add_self_flat[tensordict-eager] 0.3888ms 0.3287ms 3.0423 KOps/s 3.0142 KOps/s $\color{#35bf28}+0.93\%$
test_compile_add_self_flat[tensordict-compile] 0.5882ms 0.2515ms 3.9764 KOps/s 3.5961 KOps/s $\textbf{\color{#35bf28}+10.57\%}$
test_compile_add_self_flat[tensorclass-eager] 0.1353ms 90.6333μs 11.0335 KOps/s 10.7166 KOps/s $\color{#35bf28}+2.96\%$
test_compile_add_self_flat[tensorclass-compile] 0.2360ms 0.1864ms 5.3640 KOps/s 5.0686 KOps/s $\textbf{\color{#35bf28}+5.83\%}$
test_compile_add_self_flat[pytree-eager] 0.6290ms 0.4244ms 2.3564 KOps/s 2.3137 KOps/s $\color{#35bf28}+1.85\%$
test_compile_add_self_flat[pytree-compile] 0.2967ms 0.1929ms 5.1842 KOps/s 5.0458 KOps/s $\color{#35bf28}+2.74\%$
test_compile_copy_flat[tensordict-compile] 66.1010μs 19.7421μs 50.6531 KOps/s 50.3558 KOps/s $\color{#35bf28}+0.59\%$
test_compile_copy_flat[tensordict-eager] 84.5330μs 42.2801μs 23.6518 KOps/s 23.9002 KOps/s $\color{#d91a1a}-1.04\%$
test_compile_copy_flat[pytree-compile] 0.1180ms 21.8850μs 45.6933 KOps/s 44.8547 KOps/s $\color{#35bf28}+1.87\%$
test_compile_copy_flat[pytree-eager] 0.3270ms 52.2904μs 19.1240 KOps/s 19.0809 KOps/s $\color{#35bf28}+0.23\%$
test_compile_assign_and_add[tensordict-compile] 2.1898ms 0.2386ms 4.1911 KOps/s 3.8497 KOps/s $\textbf{\color{#35bf28}+8.87\%}$
test_compile_assign_and_add[tensordict-eager] 3.4316ms 3.2917ms 303.7921 Ops/s 301.2295 Ops/s $\color{#35bf28}+0.85\%$
test_compile_assign_and_add[pytree-compile] 2.0767ms 0.1800ms 5.5570 KOps/s 5.4072 KOps/s $\color{#35bf28}+2.77\%$
test_compile_assign_and_add[pytree-eager] 2.8904ms 2.7447ms 364.3334 Ops/s 357.0351 Ops/s $\color{#35bf28}+2.04\%$
test_compile_indexing[tensor-tensordict-compile] 0.2449ms 0.1343ms 7.4475 KOps/s 7.1855 KOps/s $\color{#35bf28}+3.65\%$
test_compile_indexing[tensor-tensordict-eager] 0.3087ms 74.1602μs 13.4843 KOps/s 12.9223 KOps/s $\color{#35bf28}+4.35\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2213ms 0.1192ms 8.3893 KOps/s 8.1256 KOps/s $\color{#35bf28}+3.25\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2510ms 44.5414μs 22.4510 KOps/s 20.6696 KOps/s $\textbf{\color{#35bf28}+8.62\%}$
test_compile_indexing[tensor-pytree-compile] 0.1679ms 0.1198ms 8.3480 KOps/s 8.1650 KOps/s $\color{#35bf28}+2.24\%$
test_compile_indexing[tensor-pytree-eager] 0.2844ms 44.3844μs 22.5304 KOps/s 20.7031 KOps/s $\textbf{\color{#35bf28}+8.83\%}$
test_compile_indexing[slice-tensordict-compile] 0.2549ms 82.0967μs 12.1808 KOps/s 11.2959 KOps/s $\textbf{\color{#35bf28}+7.83\%}$
test_compile_indexing[slice-tensordict-eager] 0.2049ms 29.1955μs 34.2519 KOps/s 35.7508 KOps/s $\color{#d91a1a}-4.19\%$
test_compile_indexing[slice-tensorclass-compile] 0.1448ms 70.4161μs 14.2013 KOps/s 13.6304 KOps/s $\color{#35bf28}+4.19\%$
test_compile_indexing[slice-tensorclass-eager] 0.2332ms 22.6474μs 44.1551 KOps/s 44.2516 KOps/s $\color{#d91a1a}-0.22\%$
test_compile_indexing[slice-pytree-compile] 0.1074ms 69.9201μs 14.3020 KOps/s 13.7956 KOps/s $\color{#35bf28}+3.67\%$
test_compile_indexing[slice-pytree-eager] 0.2491ms 22.3631μs 44.7165 KOps/s 44.5389 KOps/s $\color{#35bf28}+0.40\%$
test_compile_indexing[int-tensordict-compile] 0.1194ms 82.1634μs 12.1709 KOps/s 11.1829 KOps/s $\textbf{\color{#35bf28}+8.83\%}$
test_compile_indexing[int-tensordict-eager] 0.2243ms 27.9449μs 35.7847 KOps/s 35.5307 KOps/s $\color{#35bf28}+0.71\%$
test_compile_indexing[int-tensorclass-compile] 0.1055ms 70.0674μs 14.2720 KOps/s 13.9459 KOps/s $\color{#35bf28}+2.34\%$
test_compile_indexing[int-tensorclass-eager] 0.2458ms 22.3478μs 44.7471 KOps/s 44.4246 KOps/s $\color{#35bf28}+0.73\%$
test_compile_indexing[int-pytree-compile] 0.1177ms 70.4385μs 14.1968 KOps/s 13.7744 KOps/s $\color{#35bf28}+3.07\%$
test_compile_indexing[int-pytree-eager] 0.2441ms 22.2921μs 44.8590 KOps/s 44.2373 KOps/s $\color{#35bf28}+1.41\%$
test_compile_replace[single-eager] 84.8930μs 47.2790μs 21.1511 KOps/s 19.8702 KOps/s $\textbf{\color{#35bf28}+6.45\%}$
test_compile_replace[single-compile] 0.2275ms 0.1300ms 7.6930 KOps/s 7.2099 KOps/s $\textbf{\color{#35bf28}+6.70\%}$
test_compile_replace[multi-eager] 0.6402ms 0.5638ms 1.7736 KOps/s 1.7405 KOps/s $\color{#35bf28}+1.91\%$
test_compile_replace[multi-compile] 0.3076ms 0.1383ms 7.2285 KOps/s 7.2032 KOps/s $\color{#35bf28}+0.35\%$
test_compile_tc_getattr_20[eager] 0.2191ms 0.1627ms 6.1470 KOps/s 6.1729 KOps/s $\color{#d91a1a}-0.42\%$
test_compile_tc_getattr_20[compile] 0.3142ms 0.1479ms 6.7599 KOps/s 6.7144 KOps/s $\color{#35bf28}+0.68\%$
test_compile_clone_shallow[20-eager] 57.8310μs 19.3127μs 51.7794 KOps/s 51.0802 KOps/s $\color{#35bf28}+1.37\%$
test_compile_clone_shallow[20-compile] 73.3820μs 16.9212μs 59.0976 KOps/s 58.6251 KOps/s $\color{#35bf28}+0.81\%$
test_compile_clone_shallow[40-eager] 85.5920μs 34.0082μs 29.4047 KOps/s 28.9409 KOps/s $\color{#35bf28}+1.60\%$
test_compile_clone_shallow[40-compile] 76.9920μs 18.3248μs 54.5707 KOps/s 51.9798 KOps/s $\color{#35bf28}+4.98\%$
test_compile_clone_shallow[80-eager] 90.5930μs 63.2960μs 15.7988 KOps/s 15.4905 KOps/s $\color{#35bf28}+1.99\%$
test_compile_clone_shallow[80-compile] 93.1120μs 22.0504μs 45.3507 KOps/s 45.1246 KOps/s $\color{#35bf28}+0.50\%$
test_compile_update_inplace[eager] 0.2077ms 61.3528μs 16.2992 KOps/s 16.0965 KOps/s $\color{#35bf28}+1.26\%$
test_compile_update_inplace[compile] 0.3732ms 0.1455ms 6.8710 KOps/s 6.7379 KOps/s $\color{#35bf28}+1.98\%$
test_mod_add[eager] 0.2786ms 0.2280ms 4.3858 KOps/s 4.4141 KOps/s $\color{#d91a1a}-0.64\%$
test_mod_add[compile] 0.2003ms 0.1312ms 7.6208 KOps/s 7.3171 KOps/s $\color{#35bf28}+4.15\%$
test_mod_add[compile-overhead] 0.3425ms 0.1771ms 5.6453 KOps/s 5.4455 KOps/s $\color{#35bf28}+3.67\%$
test_mod_wrap[eager] 0.5563ms 0.4698ms 2.1285 KOps/s 2.1151 KOps/s $\color{#35bf28}+0.63\%$
test_mod_wrap[compile] 0.4890ms 0.3943ms 2.5361 KOps/s 2.4364 KOps/s $\color{#35bf28}+4.09\%$
test_mod_wrap[compile-overhead] 8.4833ms 4.6560ms 214.7752 Ops/s 214.1374 Ops/s $\color{#35bf28}+0.30\%$
test_mod_wrap_and_backward[eager] 1.9462ms 1.7873ms 559.5018 Ops/s 582.4839 Ops/s $\color{#d91a1a}-3.95\%$
test_mod_wrap_and_backward[compile] 2.4719ms 1.5866ms 630.2964 Ops/s 634.1531 Ops/s $\color{#d91a1a}-0.61\%$
test_mod_wrap_and_backward[compile-overhead] 1.6900ms 0.9406ms 1.0631 KOps/s 1.0359 KOps/s $\color{#35bf28}+2.62\%$
test_seq_add[eager] 0.6170ms 0.5164ms 1.9363 KOps/s 1.9268 KOps/s $\color{#35bf28}+0.49\%$
test_seq_add[compile] 0.2514ms 0.1403ms 7.1278 KOps/s 6.6249 KOps/s $\textbf{\color{#35bf28}+7.59\%}$
test_seq_add[compile-overhead] 0.2375ms 0.1825ms 5.4790 KOps/s 5.1575 KOps/s $\textbf{\color{#35bf28}+6.23\%}$
test_seq_wrap[eager] 1.5268ms 1.4196ms 704.4068 Ops/s 696.9301 Ops/s $\color{#35bf28}+1.07\%$
test_seq_wrap[compile] 0.5141ms 0.4083ms 2.4491 KOps/s 2.3592 KOps/s $\color{#35bf28}+3.81\%$
test_seq_wrap[compile-overhead] 0.3734ms 0.3032ms 3.2986 KOps/s 3.2058 KOps/s $\color{#35bf28}+2.89\%$
test_func_call_runtime[False-eager] 0.8673ms 0.7975ms 1.2540 KOps/s 1.2115 KOps/s $\color{#35bf28}+3.51\%$
test_func_call_runtime[False-compile] 1.0710ms 0.9790ms 1.0215 KOps/s 963.7391 Ops/s $\textbf{\color{#35bf28}+5.99\%}$
test_func_call_runtime[False-compile-overhead] 0.6587ms 0.5114ms 1.9554 KOps/s 1.9312 KOps/s $\color{#35bf28}+1.25\%$
test_func_call_runtime[True-eager] 1.1371ms 1.0398ms 961.7100 Ops/s 948.7969 Ops/s $\color{#35bf28}+1.36\%$
test_func_call_runtime[True-compile] 1.1209ms 0.9898ms 1.0103 KOps/s 994.3443 Ops/s $\color{#35bf28}+1.60\%$
test_func_call_runtime[True-compile-overhead] 0.6054ms 0.5268ms 1.8983 KOps/s 1.8741 KOps/s $\color{#35bf28}+1.29\%$
test_func_call_cm_runtime[False-eager] 0.8911ms 0.7886ms 1.2680 KOps/s 1.2502 KOps/s $\color{#35bf28}+1.43\%$
test_func_call_cm_runtime[False-compile] 1.0692ms 0.9604ms 1.0412 KOps/s 1.0238 KOps/s $\color{#35bf28}+1.70\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5442ms 0.4937ms 2.0255 KOps/s 2.0041 KOps/s $\color{#35bf28}+1.07\%$
test_func_call_cm_runtime[True-eager] 1.3714ms 1.1887ms 841.2901 Ops/s 834.5445 Ops/s $\color{#35bf28}+0.81\%$
test_func_call_cm_runtime[True-compile] 1.1191ms 1.0160ms 984.2067 Ops/s 967.3396 Ops/s $\color{#35bf28}+1.74\%$
test_func_call_cm_runtime[True-compile-overhead] 0.7025ms 0.5498ms 1.8190 KOps/s 1.7829 KOps/s $\color{#35bf28}+2.02\%$
test_vmap_func_call_cm_runtime[eager] 2.8042ms 2.3609ms 423.5591 Ops/s 420.2703 Ops/s $\color{#35bf28}+0.78\%$
test_vmap_func_call_cm_runtime[compile] 1.1277ms 1.0197ms 980.6542 Ops/s 971.1880 Ops/s $\color{#35bf28}+0.97\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.5967ms 0.5390ms 1.8551 KOps/s 1.8316 KOps/s $\color{#35bf28}+1.29\%$
test_distributed 0.7352ms 0.1499ms 6.6713 KOps/s 6.5157 KOps/s $\color{#35bf28}+2.39\%$
test_tdmodule 0.5874ms 0.1948ms 5.1323 KOps/s 5.0104 KOps/s $\color{#35bf28}+2.43\%$
test_tdmodule_dispatch 0.4284ms 0.3725ms 2.6847 KOps/s 2.6434 KOps/s $\color{#35bf28}+1.56\%$
test_tdseq 45.6310μs 26.4574μs 37.7966 KOps/s 36.4688 KOps/s $\color{#35bf28}+3.64\%$
test_tdseq_dispatch 0.2594ms 0.2146ms 4.6605 KOps/s 4.5747 KOps/s $\color{#35bf28}+1.88\%$
test_instantiation_functorch 2.1891ms 2.0818ms 480.3474 Ops/s 479.4287 Ops/s $\color{#35bf28}+0.19\%$
test_exec_functorch 0.2255ms 0.1836ms 5.4473 KOps/s 5.5044 KOps/s $\color{#d91a1a}-1.04\%$
test_exec_functional_call 0.2254ms 0.1633ms 6.1244 KOps/s 5.9632 KOps/s $\color{#35bf28}+2.70\%$
test_exec_td_decorator 0.4565ms 0.2404ms 4.1599 KOps/s 4.0912 KOps/s $\color{#35bf28}+1.68\%$
test_vmap_mlp_speed_decorator[True-True] 1.2144ms 1.0207ms 979.7524 Ops/s 985.1431 Ops/s $\color{#d91a1a}-0.55\%$
test_vmap_mlp_speed_decorator[True-False] 1.2452ms 1.0200ms 980.4046 Ops/s 987.3814 Ops/s $\color{#d91a1a}-0.71\%$
test_vmap_mlp_speed_decorator[False-True] 0.9151ms 0.7193ms 1.3903 KOps/s 1.3634 KOps/s $\color{#35bf28}+1.97\%$
test_vmap_mlp_speed_decorator[False-False] 0.8927ms 0.7182ms 1.3924 KOps/s 1.3617 KOps/s $\color{#35bf28}+2.26\%$
test_vmap_transformer_speed_decorator[True-True] 21.1088ms 20.8663ms 47.9241 Ops/s 47.8533 Ops/s $\color{#35bf28}+0.15\%$
test_vmap_transformer_speed_decorator[True-False] 21.4637ms 20.8803ms 47.8920 Ops/s 47.9599 Ops/s $\color{#d91a1a}-0.14\%$
test_vmap_transformer_speed_decorator[False-True] 21.3354ms 20.4992ms 48.7824 Ops/s 48.8273 Ops/s $\color{#d91a1a}-0.09\%$
test_vmap_transformer_speed_decorator[False-False] 21.1831ms 20.4775ms 48.8341 Ops/s 48.8587 Ops/s $\color{#d91a1a}-0.05\%$
test_to_module_speed[True] 2.1486ms 1.5021ms 665.7201 Ops/s 672.2634 Ops/s $\color{#d91a1a}-0.97\%$
test_to_module_speed[False] 2.0889ms 1.4653ms 682.4764 Ops/s 692.2564 Ops/s $\color{#d91a1a}-1.41\%$
test_tc_init 76.4620μs 42.8574μs 23.3332 KOps/s 22.9234 KOps/s $\color{#35bf28}+1.79\%$
test_tc_init_tensor_only 30.5200μs 9.0381μs 110.6426 KOps/s 109.7744 KOps/s $\color{#35bf28}+0.79\%$
test_tc_init_nested 0.1326ms 85.4352μs 11.7048 KOps/s 11.5293 KOps/s $\color{#35bf28}+1.52\%$
test_tc_init_many_fields 33.7410μs 14.8285μs 67.4375 KOps/s 67.9262 KOps/s $\color{#d91a1a}-0.72\%$
test_tc_first_layer_tensor 20.8800μs 1.8226μs 548.6678 KOps/s 550.6861 KOps/s $\color{#d91a1a}-0.37\%$
test_tc_first_layer_tensor_only 2.3476μs 0.3915μs 2.5545 MOps/s 2.5292 MOps/s $\color{#35bf28}+1.00\%$
test_tc_first_layer_tensor_set 88.2030μs 3.9199μs 255.1075 KOps/s 253.2944 KOps/s $\color{#35bf28}+0.72\%$
test_tc_first_layer_tensor_only_set 18.8400μs 3.2495μs 307.7357 KOps/s 303.3562 KOps/s $\color{#35bf28}+1.44\%$
test_tc_first_layer_nontensor 35.8110μs 6.1725μs 162.0102 KOps/s 156.2797 KOps/s $\color{#35bf28}+3.67\%$
test_tc_second_layer_tensor 37.9310μs 4.4067μs 226.9296 KOps/s 223.7751 KOps/s $\color{#35bf28}+1.41\%$
test_tc_second_layer_nontensor 32.0010μs 8.7430μs 114.3772 KOps/s 112.3737 KOps/s $\color{#35bf28}+1.78\%$
test_unbind 0.2852s 16.5852ms 60.2947 Ops/s 71.8893 Ops/s $\textbf{\color{#d91a1a}-16.13\%}$
test_full_like 5.8284ms 4.3681ms 228.9300 Ops/s 230.5770 Ops/s $\color{#d91a1a}-0.71\%$
test_zeros_like 4.8381ms 4.3633ms 229.1848 Ops/s 233.0004 Ops/s $\color{#d91a1a}-1.64\%$
test_ones_like 4.4522ms 4.3301ms 230.9417 Ops/s 232.4303 Ops/s $\color{#d91a1a}-0.64\%$
test_clone 6.5936ms 6.4157ms 155.8679 Ops/s 156.1175 Ops/s $\color{#d91a1a}-0.16\%$
test_squeeze 66.4210μs 14.6362μs 68.3239 KOps/s 67.8097 KOps/s $\color{#35bf28}+0.76\%$
test_unsqueeze 0.1599ms 0.1116ms 8.9598 KOps/s 9.0417 KOps/s $\color{#d91a1a}-0.91\%$
test_split 0.3842ms 0.1825ms 5.4781 KOps/s 5.5101 KOps/s $\color{#d91a1a}-0.58\%$
test_permute 0.2804ms 0.2052ms 4.8741 KOps/s 4.8449 KOps/s $\color{#35bf28}+0.60\%$
test_stack 52.2629ms 51.4415ms 19.4395 Ops/s 19.4837 Ops/s $\color{#d91a1a}-0.23\%$
test_cat 51.7615ms 51.3960ms 19.4568 Ops/s 19.5283 Ops/s $\color{#d91a1a}-0.37\%$
test_sequential_tensordict 0.7273ms 0.5740ms 1.7423 KOps/s 1.7451 KOps/s $\color{#d91a1a}-0.16\%$
test_sequential_graph_module 0.5309ms 0.1174ms 8.5159 KOps/s 8.1033 KOps/s $\textbf{\color{#35bf28}+5.09\%}$
test_nested_tensordict 1.7577ms 0.8237ms 1.2140 KOps/s 1.2231 KOps/s $\color{#d91a1a}-0.75\%$
test_nested_graph_module 0.1858ms 0.1346ms 7.4319 KOps/s 7.6327 KOps/s $\color{#d91a1a}-2.63\%$

@vmoens
Update
7517c6c 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request May 24, 2026
@vmoens @cursoragent
[Compile] Locked-fast-path for clone(recurse=False) via cached schema
78b8785 
A locked ``TensorDict`` already guarantees that its top-level key set
is fixed. Cache that key set in a ``_LockedSchema`` namedtuple when
``lock_()`` is called and consume it in the hot path so Dynamo can
guard on an immutable tuple instead of walking the backing dict.

Concretely, under ``torch.compile`` ``clone(recurse=False)`` on a
locked TD switches from

    result._tensordict.update(self._tensordict)

(which forces a ``DICT_LENGTH`` / dict-keys guard on ``_tensordict``)
to a per-key copy driven by the cached ``schema.keys`` tuple. Cheap
``LENGTH_CHECK`` / ``EQUALS_MATCH`` guards on the immutable schema
replace the dict-internals guard, and the schema can be shared
across locked TDs that lock to the same key set.

Plumbing:

- ``_LockedSchema`` namedtuple added in ``tensordict.utils``.
- ``TensorDictBase._locked_schema`` class default + new
  ``_build_lock_schema()`` hook called from ``_propagate_lock`` in
  eager mode (skipped under compile to avoid the very guard we want
  to avoid building).
- ``_propagate_unlock`` drops the schema so it cannot go stale once
  keys can be added/removed again.
- ``TensorDict._build_lock_schema`` snapshots ``tuple(self._tensordict)``.
- ``TensorDict._clone`` uses the cached schema under
  ``is_compiling() and self._is_locked``.

Adds guard-count tests that exercise the locked clone fast path and
verify the schema is cleared on unlock and rebuilt on relock.

The locked-fast-path infrastructure is intentionally narrow in this
commit: only ``clone(recurse=False)`` consumes the schema. Extending
it to ``_set_str`` / ``_has_names`` is left for follow-ups so the
schema-cache invariants can be evaluated under real workloads first.

Authored with Claude.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: ff20e12
Pull-Request: #1715
vmoens added a commit that referenced this pull request May 24, 2026
@vmoens @cursoragent
[Compile] Locked-fast-path for clone(recurse=False) via cached schema
445f1a3 
A locked ``TensorDict`` already guarantees that its top-level key set
is fixed. Cache that key set in a ``_LockedSchema`` namedtuple when
``lock_()`` is called and consume it in the hot path so Dynamo can
guard on an immutable tuple instead of walking the backing dict.

Concretely, under ``torch.compile`` ``clone(recurse=False)`` on a
locked TD switches from

    result._tensordict.update(self._tensordict)

(which forces a ``DICT_LENGTH`` / dict-keys guard on ``_tensordict``)
to a per-key copy driven by the cached ``schema.keys`` tuple. Cheap
``LENGTH_CHECK`` / ``EQUALS_MATCH`` guards on the immutable schema
replace the dict-internals guard, and the schema can be shared
across locked TDs that lock to the same key set.

Plumbing:

- ``_LockedSchema`` namedtuple added in ``tensordict.utils``.
- ``TensorDictBase._locked_schema`` class default + new
  ``_build_lock_schema()`` hook called from ``_propagate_lock`` in
  eager mode (skipped under compile to avoid the very guard we want
  to avoid building).
- ``_propagate_unlock`` drops the schema so it cannot go stale once
  keys can be added/removed again.
- ``TensorDict._build_lock_schema`` snapshots ``tuple(self._tensordict)``.
- ``TensorDict._clone`` uses the cached schema under
  ``is_compiling() and self._is_locked``.

Adds guard-count tests that exercise the locked clone fast path and
verify the schema is cleared on unlock and rebuilt on relock.

The locked-fast-path infrastructure is intentionally narrow in this
commit: only ``clone(recurse=False)`` consumes the schema. Extending
it to ``_set_str`` / ``_has_names`` is left for follow-ups so the
schema-cache invariants can be evaluated under real workloads first.

Authored with Claude.

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: ff20e12
Pull-Request: #1715
@vmoens
vmoens merged commit 7517c6c into gh/vmoens/95/base May 24, 2026
38 of 44 checks passed
@vmoens
vmoens deleted the gh/vmoens/95/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