Skip to content

[BugFix] Preserve tensorclass identity in memmap_/load_memmap for TensorClass subclasses - #1722

Merged
vmoens merged 1 commit into
mainfrom
fix-memmap-tensorclass
Jun 11, 2026
Merged

[BugFix] Preserve tensorclass identity in memmap_/load_memmap for TensorClass subclasses#1722
vmoens merged 1 commit into
mainfrom
fix-memmap-tensorclass

Conversation

@vmoens

@vmoens vmoens commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Description

A tensorclass written with the subclass syntax (class Foo(TensorClass["nocast"])) did not survive a memmap_ / load_memmap roundtrip: it came back as a plain TensorDict holding the tensorclass fields, losing the class identity. This affected both top-level and nested instances. Decorator-based tensorclasses (@tensorclass) were unaffected at the top level, but their subclasses (@tensorclass class Child(Parent)) hit the same problem on load.

import tempfile, torch
from tensordict import TensorDict, TensorClass

class Foo(TensorClass["nocast"]):
    a: torch.Tensor
    meta: str | None = None

foo = Foo(a=torch.zeros(5), meta="hello", batch_size=[5])
td = TensorDict({"obs": {"foo": foo}}, batch_size=[5])
with tempfile.TemporaryDirectory() as tmp:
    td.memmap_(tmp)
    back = TensorDict.load_memmap(tmp)
    print(type(back["obs", "foo"]))  # was: TensorDict -- now: Foo

Root cause

Two wiring bugs in tensordict/tensorclass.py:

  1. Save side: "_memmap_" was listed in _FALLBACK_METHOD_FROM_TD (added in the bulk registration of [Refactor] Add missing functions in tensorclass register #1153). In _patch_tc, the fallback loop runs after the explicit cls._memmap_ = _memmap_ assignment and overwrote it with a wrapper delegating to self._tensordict._memmap_(...). The node was therefore saved as a plain TensorDict (flat fields, _type: TensorDict, no _tensordict/ subdirectory), so load_memmap had nothing to dispatch on. The dedicated tensorclass _memmap_, which records the class in meta.json, is now preserved.

  2. Load side: the classmethod-wrapping loop in _tensorclass only checked cls.__dict__ when deciding whether a TensorDict classmethod was "missing". For TensorClass subclasses, the genuine _load_memmap classmethod is inherited from the TensorClass base, so the loop shadowed it with TensorDict._load_memmap, which silently returns an empty instance on the tensorclass directory layout. The loop now skips attributes already provided by an ancestor, unless they are _wrap_classmethod wrappers, which hardcode the parent class and must be re-wrapped per subclass so that e.g. Child.from_json(...) still returns a Child. (The NonTensorData classes had been dodging this bug only by re-assigning _load_memmap = classmethod(_load_memmap) in their class bodies.)

Backward compatibility

Directories saved with the old (buggy) flat layout still load without error, degrading to a plain TensorDict exactly as before.

Tests

Two tests added to TestMemmap in test/test_tensorclass.py:

  • test_memmap_nested_tensorclass_roundtrip: both tensorclass flavors nested at depth >= 2 inside a plain TensorDict, including a non-tensor field, checking class identity and values.
  • test_memmap_tensorclass_subclass_roundtrip: top-level TensorClass subclass through both TensorDict.load_memmap and cls.load_memmap.

Also verified: non-inplace memmap(), memmap_(num_threads=4), memmap_like(), decorator-subclass roundtrips, and the existing suites (test_tensorclass.py, test_memmap.py, test_td_memmap_split.py, test_compatibility.py, test_typedtensordict.py) pass.

🤖 Generated with Claude Code

@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 Jun 10, 2026
@github-actions github-actions Bot added bug Something isn't working CI Test tensorclass labels Jun 10, 2026
@github-actions

github-actions Bot commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

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

Total Benchmarks: 344. Improved: $\large\color{#35bf28}10$. 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.2103ms 0.1680ms 5.9538 KOps/s 5.9584 KOps/s $\color{#d91a1a}-0.08\%$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.4382ms 0.2982ms 3.3537 KOps/s 3.3656 KOps/s $\color{#d91a1a}-0.35\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8802ms 0.7749ms 1.2905 KOps/s 1.2738 KOps/s $\color{#35bf28}+1.31\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.3752ms 0.3403ms 2.9388 KOps/s 2.8771 KOps/s $\color{#35bf28}+2.15\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.4775ms 1.1906ms 839.9127 Ops/s 898.1631 Ops/s $\textbf{\color{#d91a1a}-6.49\%}$
test_collector_preallocated_write[stack_out-cpu-4096-32] 8.0665ms 7.6818ms 130.1786 Ops/s 130.4911 Ops/s $\color{#d91a1a}-0.24\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.2938ms 0.2459ms 4.0673 KOps/s 4.0450 KOps/s $\color{#35bf28}+0.55\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.5106ms 0.3694ms 2.7071 KOps/s 2.6936 KOps/s $\color{#35bf28}+0.50\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.9400ms 0.8196ms 1.2202 KOps/s 1.2166 KOps/s $\color{#35bf28}+0.29\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.3421ms 0.2486ms 4.0222 KOps/s 3.8797 KOps/s $\color{#35bf28}+3.67\%$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.5110ms 0.3841ms 2.6035 KOps/s 2.6797 KOps/s $\color{#d91a1a}-2.84\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 0.9868ms 0.8888ms 1.1251 KOps/s 1.1334 KOps/s $\color{#d91a1a}-0.74\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.2880ms 0.2148ms 4.6563 KOps/s 4.5299 KOps/s $\color{#35bf28}+2.79\%$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 0.9690ms 0.8346ms 1.1981 KOps/s 1.1748 KOps/s $\color{#35bf28}+1.98\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.4533ms 3.3095ms 302.1596 Ops/s 295.3309 Ops/s $\color{#35bf28}+2.31\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.4373ms 0.3688ms 2.7114 KOps/s 2.6714 KOps/s $\color{#35bf28}+1.50\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 2.0507ms 1.8398ms 543.5314 Ops/s 582.4147 Ops/s $\textbf{\color{#d91a1a}-6.68\%}$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 10.9315ms 10.7086ms 93.3828 Ops/s 93.6327 Ops/s $\color{#d91a1a}-0.27\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.4362ms 0.3917ms 2.5527 KOps/s 2.5315 KOps/s $\color{#35bf28}+0.84\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.7168ms 1.4978ms 667.6681 Ops/s 682.9675 Ops/s $\color{#d91a1a}-2.24\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 6.1386ms 5.8683ms 170.4071 Ops/s 174.1999 Ops/s $\color{#d91a1a}-2.18\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.4798ms 0.4092ms 2.4440 KOps/s 2.5483 KOps/s $\color{#d91a1a}-4.09\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.6458ms 1.5482ms 645.9165 Ops/s 677.9056 Ops/s $\color{#d91a1a}-4.72\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 6.4050ms 5.9430ms 168.2648 Ops/s 172.4022 Ops/s $\color{#d91a1a}-2.40\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.6016ms 0.1730ms 5.7791 KOps/s 5.7923 KOps/s $\color{#d91a1a}-0.23\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 0.7560ms 0.6713ms 1.4896 KOps/s 1.4787 KOps/s $\color{#35bf28}+0.74\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 2.8744ms 2.6814ms 372.9410 Ops/s 373.0676 Ops/s $\color{#d91a1a}-0.03\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.3949ms 0.3155ms 3.1696 KOps/s 3.0782 KOps/s $\color{#35bf28}+2.97\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 1.8145ms 1.6170ms 618.4281 Ops/s 667.9433 Ops/s $\textbf{\color{#d91a1a}-7.41\%}$
test_collector_preallocated_write[copy_at-cpu-4096-32] 11.4411ms 10.1933ms 98.1037 Ops/s 100.1456 Ops/s $\color{#d91a1a}-2.04\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.4292ms 0.3201ms 3.1237 KOps/s 3.0941 KOps/s $\color{#35bf28}+0.96\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.3962ms 1.2133ms 824.2095 Ops/s 834.2521 Ops/s $\color{#d91a1a}-1.20\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 5.0199ms 4.7369ms 211.1101 Ops/s 214.1224 Ops/s $\color{#d91a1a}-1.41\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.3910ms 0.3211ms 3.1142 KOps/s 3.0887 KOps/s $\color{#35bf28}+0.83\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.3178ms 1.1944ms 837.2660 Ops/s 831.1108 Ops/s $\color{#35bf28}+0.74\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 5.0764ms 4.6691ms 214.1752 Ops/s 212.0092 Ops/s $\color{#35bf28}+1.02\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1607ms 0.1217ms 8.2147 KOps/s 8.1559 KOps/s $\color{#35bf28}+0.72\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5442ms 0.4835ms 2.0684 KOps/s 2.0770 KOps/s $\color{#d91a1a}-0.42\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 1.9968ms 1.9140ms 522.4779 Ops/s 523.1903 Ops/s $\color{#d91a1a}-0.14\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.4303ms 0.2720ms 3.6764 KOps/s 3.8724 KOps/s $\textbf{\color{#d91a1a}-5.06\%}$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.6524ms 1.4550ms 687.2639 Ops/s 783.0103 Ops/s $\textbf{\color{#d91a1a}-12.23\%}$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 9.2586ms 8.9167ms 112.1496 Ops/s 111.4219 Ops/s $\color{#35bf28}+0.65\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.3607ms 0.3009ms 3.3230 KOps/s 3.6204 KOps/s $\textbf{\color{#d91a1a}-8.22\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.2725ms 1.1103ms 900.6861 Ops/s 971.4328 Ops/s $\textbf{\color{#d91a1a}-7.28\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.4852ms 4.3348ms 230.6907 Ops/s 247.5589 Ops/s $\textbf{\color{#d91a1a}-6.81\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.3864ms 0.3033ms 3.2972 KOps/s 3.6156 KOps/s $\textbf{\color{#d91a1a}-8.81\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.2155ms 1.1180ms 894.4487 Ops/s 967.5217 Ops/s $\textbf{\color{#d91a1a}-7.55\%}$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.5294ms 4.3940ms 227.5816 Ops/s 247.1965 Ops/s $\textbf{\color{#d91a1a}-7.93\%}$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 17.1442ms 16.5817ms 60.3073 Ops/s 60.8398 Ops/s $\color{#d91a1a}-0.88\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 33.6009ms 33.2378ms 30.0862 Ops/s 30.1831 Ops/s $\color{#d91a1a}-0.32\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 89.9042ms 89.4891ms 11.1745 Ops/s 10.9424 Ops/s $\color{#35bf28}+2.12\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2119s 0.2111s 4.7365 Ops/s 4.7510 Ops/s $\color{#d91a1a}-0.30\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 13.9872ms 13.7442ms 72.7581 Ops/s 72.5333 Ops/s $\color{#35bf28}+0.31\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 27.3446ms 27.1448ms 36.8394 Ops/s 36.7350 Ops/s $\color{#35bf28}+0.28\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 14.0893ms 13.8849ms 72.0206 Ops/s 71.0445 Ops/s $\color{#35bf28}+1.37\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 27.6383ms 27.4128ms 36.4794 Ops/s 36.0715 Ops/s $\color{#35bf28}+1.13\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 55.9095ms 55.4214ms 18.0436 Ops/s 18.8414 Ops/s $\color{#d91a1a}-4.23\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1119s 0.1091s 9.1619 Ops/s 9.4064 Ops/s $\color{#d91a1a}-2.60\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1593s 0.1557s 6.4235 Ops/s 6.6197 Ops/s $\color{#d91a1a}-2.97\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3473s 0.3441s 2.9057 Ops/s 2.9460 Ops/s $\color{#d91a1a}-1.37\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 80.8749ms 80.1098ms 12.4829 Ops/s 13.3982 Ops/s $\textbf{\color{#d91a1a}-6.83\%}$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1614s 0.1586s 6.3068 Ops/s 6.6914 Ops/s $\textbf{\color{#d91a1a}-5.75\%}$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 80.6071ms 77.7876ms 12.8555 Ops/s 13.2702 Ops/s $\color{#d91a1a}-3.12\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1650s 0.1621s 6.1700 Ops/s 6.6285 Ops/s $\textbf{\color{#d91a1a}-6.92\%}$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.8158ms 30.3241ms 32.9771 Ops/s 32.4338 Ops/s $\color{#35bf28}+1.68\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 62.1214ms 61.6837ms 16.2117 Ops/s 16.1246 Ops/s $\color{#35bf28}+0.54\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1292s 0.1270s 7.8739 Ops/s 7.9910 Ops/s $\color{#d91a1a}-1.46\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2935s 0.2894s 3.4553 Ops/s 3.5856 Ops/s $\color{#d91a1a}-3.64\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 43.5424ms 42.6401ms 23.4521 Ops/s 23.4356 Ops/s $\color{#35bf28}+0.07\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 85.5460ms 84.8769ms 11.7818 Ops/s 11.6948 Ops/s $\color{#35bf28}+0.74\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 43.3157ms 42.6988ms 23.4199 Ops/s 23.1539 Ops/s $\color{#35bf28}+1.15\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 86.1179ms 85.4083ms 11.7085 Ops/s 11.6212 Ops/s $\color{#35bf28}+0.75\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 20.7673ms 20.5457ms 48.6721 Ops/s 47.8808 Ops/s $\color{#35bf28}+1.65\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 42.5667ms 42.1360ms 23.7327 Ops/s 23.3538 Ops/s $\color{#35bf28}+1.62\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1189s 0.1173s 8.5279 Ops/s 8.7866 Ops/s $\color{#d91a1a}-2.94\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2713s 0.2671s 3.7445 Ops/s 3.8742 Ops/s $\color{#d91a1a}-3.35\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 37.1570ms 36.3423ms 27.5161 Ops/s 27.1075 Ops/s $\color{#35bf28}+1.51\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 73.4650ms 72.8646ms 13.7241 Ops/s 13.5048 Ops/s $\color{#35bf28}+1.62\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 37.0375ms 36.2150ms 27.6129 Ops/s 26.9071 Ops/s $\color{#35bf28}+2.62\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 73.2010ms 72.4469ms 13.8032 Ops/s 13.4548 Ops/s $\color{#35bf28}+2.59\%$
test_plain_set_nested 56.9610μs 15.5163μs 64.4482 KOps/s 64.7107 KOps/s $\color{#d91a1a}-0.41\%$
test_plain_set_stack_nested 39.6510μs 15.3113μs 65.3113 KOps/s 66.2982 KOps/s $\color{#d91a1a}-1.49\%$
test_plain_set_nested_inplace 44.2910μs 16.8662μs 59.2903 KOps/s 59.5934 KOps/s $\color{#d91a1a}-0.51\%$
test_plain_set_stack_nested_inplace 45.5210μs 16.4753μs 60.6970 KOps/s 60.4785 KOps/s $\color{#35bf28}+0.36\%$
test_items 31.7400μs 6.0459μs 165.4011 KOps/s 166.8352 KOps/s $\color{#d91a1a}-0.86\%$
test_items_nested 0.5555ms 0.4581ms 2.1832 KOps/s 2.1769 KOps/s $\color{#35bf28}+0.29\%$
test_items_nested_locked 0.5617ms 0.4621ms 2.1642 KOps/s 2.1495 KOps/s $\color{#35bf28}+0.69\%$
test_items_nested_leaf 0.1292ms 97.4511μs 10.2616 KOps/s 10.3209 KOps/s $\color{#d91a1a}-0.58\%$
test_items_stack_nested 0.5844ms 0.4601ms 2.1734 KOps/s 2.1732 KOps/s $+0.01\%$
test_items_stack_nested_leaf 0.1853ms 96.4156μs 10.3718 KOps/s 10.4043 KOps/s $\color{#d91a1a}-0.31\%$
test_items_stack_nested_locked 0.5518ms 0.4638ms 2.1563 KOps/s 2.1810 KOps/s $\color{#d91a1a}-1.13\%$
test_keys 27.4110μs 4.1810μs 239.1796 KOps/s 228.2644 KOps/s $\color{#35bf28}+4.78\%$
test_keys_nested 0.1868ms 0.1294ms 7.7298 KOps/s 7.6331 KOps/s $\color{#35bf28}+1.27\%$
test_keys_nested_locked 2.1964ms 0.1381ms 7.2432 KOps/s 7.1323 KOps/s $\color{#35bf28}+1.56\%$
test_keys_nested_leaf 0.1922ms 0.1205ms 8.2957 KOps/s 8.2619 KOps/s $\color{#35bf28}+0.41\%$
test_keys_stack_nested 0.1797ms 0.1287ms 7.7696 KOps/s 7.6139 KOps/s $\color{#35bf28}+2.04\%$
test_keys_stack_nested_leaf 0.1751ms 0.1197ms 8.3523 KOps/s 8.2924 KOps/s $\color{#35bf28}+0.72\%$
test_keys_stack_nested_locked 0.1895ms 0.1367ms 7.3180 KOps/s 7.2074 KOps/s $\color{#35bf28}+1.53\%$
test_values 15.0144μs 1.0220μs 978.4625 KOps/s 983.6465 KOps/s $\color{#d91a1a}-0.53\%$
test_values_nested 87.1420μs 52.4652μs 19.0602 KOps/s 18.7774 KOps/s $\color{#35bf28}+1.51\%$
test_values_nested_locked 0.1157ms 55.7600μs 17.9340 KOps/s 17.6578 KOps/s $\color{#35bf28}+1.56\%$
test_values_nested_leaf 0.1124ms 60.1583μs 16.6228 KOps/s 16.4037 KOps/s $\color{#35bf28}+1.34\%$
test_values_stack_nested 88.9320μs 52.4457μs 19.0674 KOps/s 18.8555 KOps/s $\color{#35bf28}+1.12\%$
test_values_stack_nested_leaf 0.1361ms 60.5243μs 16.5223 KOps/s 16.3918 KOps/s $\color{#35bf28}+0.80\%$
test_values_stack_nested_locked 79.1630μs 55.9585μs 17.8704 KOps/s 17.5159 KOps/s $\color{#35bf28}+2.02\%$
test_membership 5.9035μs 0.8284μs 1.2071 MOps/s 1.1774 MOps/s $\color{#35bf28}+2.52\%$
test_membership_nested 33.3210μs 2.9512μs 338.8414 KOps/s 331.2453 KOps/s $\color{#35bf28}+2.29\%$
test_membership_nested_leaf 43.4710μs 2.9493μs 339.0683 KOps/s 330.2059 KOps/s $\color{#35bf28}+2.68\%$
test_membership_stacked_nested 76.3720μs 2.8478μs 351.1494 KOps/s 342.7730 KOps/s $\color{#35bf28}+2.44\%$
test_membership_stacked_nested_leaf 41.5410μs 2.8624μs 349.3532 KOps/s 343.0195 KOps/s $\color{#35bf28}+1.85\%$
test_membership_nested_last 33.9410μs 4.4385μs 225.3020 KOps/s 222.5775 KOps/s $\color{#35bf28}+1.22\%$
test_membership_nested_leaf_last 38.9210μs 4.4789μs 223.2693 KOps/s 222.0658 KOps/s $\color{#35bf28}+0.54\%$
test_membership_stacked_nested_last 35.7710μs 4.3385μs 230.4968 KOps/s 227.7187 KOps/s $\color{#35bf28}+1.22\%$
test_membership_stacked_nested_leaf_last 45.2220μs 4.3356μs 230.6508 KOps/s 227.4961 KOps/s $\color{#35bf28}+1.39\%$
test_nested_getleaf 45.6520μs 21.8064μs 45.8582 KOps/s 46.6324 KOps/s $\color{#d91a1a}-1.66\%$
test_nested_get 0.1028ms 20.5656μs 48.6250 KOps/s 48.4900 KOps/s $\color{#35bf28}+0.28\%$
test_stacked_getleaf 53.5820μs 21.7299μs 46.0196 KOps/s 46.8585 KOps/s $\color{#d91a1a}-1.79\%$
test_stacked_get 43.7620μs 20.5496μs 48.6628 KOps/s 48.8279 KOps/s $\color{#d91a1a}-0.34\%$
test_nested_getitemleaf 57.4710μs 22.1103μs 45.2277 KOps/s 44.9951 KOps/s $\color{#35bf28}+0.52\%$
test_nested_getitem 48.2410μs 20.9706μs 47.6858 KOps/s 47.2443 KOps/s $\color{#35bf28}+0.93\%$
test_stacked_getitemleaf 90.8120μs 21.9550μs 45.5476 KOps/s 45.5282 KOps/s $\color{#35bf28}+0.04\%$
test_stacked_getitem 59.3310μs 20.7069μs 48.2930 KOps/s 48.1410 KOps/s $\color{#35bf28}+0.32\%$
test_lock_nested 0.5700ms 0.4944ms 2.0225 KOps/s 1.9763 KOps/s $\color{#35bf28}+2.34\%$
test_lock_stack_nested 0.5771ms 0.5041ms 1.9837 KOps/s 1.9816 KOps/s $\color{#35bf28}+0.11\%$
test_unlock_nested 0.4646ms 0.3928ms 2.5458 KOps/s 2.5461 KOps/s $\color{#d91a1a}-0.01\%$
test_unlock_stack_nested 0.4724ms 0.3968ms 2.5201 KOps/s 2.5124 KOps/s $\color{#35bf28}+0.31\%$
test_flatten_speed 0.1588ms 0.1233ms 8.1074 KOps/s 8.1768 KOps/s $\color{#d91a1a}-0.85\%$
test_unflatten_speed 0.6543ms 0.5692ms 1.7569 KOps/s 1.7688 KOps/s $\color{#d91a1a}-0.67\%$
test_common_ops 0.1037s 1.7999ms 555.5927 Ops/s 269.8616 Ops/s $\textbf{\color{#35bf28}+105.88\%}$
test_creation 0.1194ms 3.1424μs 318.2305 KOps/s 316.3031 KOps/s $\color{#35bf28}+0.61\%$
test_creation_empty 55.8010μs 7.0318μs 142.2111 KOps/s 141.8686 KOps/s $\color{#35bf28}+0.24\%$
test_creation_nested_1 42.6210μs 11.5355μs 86.6888 KOps/s 86.6710 KOps/s $\color{#35bf28}+0.02\%$
test_creation_nested_2 51.5010μs 13.5536μs 73.7809 KOps/s 73.3717 KOps/s $\color{#35bf28}+0.56\%$
test_creation_many_keys[10] 47.6110μs 21.1593μs 47.2604 KOps/s 47.0059 KOps/s $\color{#35bf28}+0.54\%$
test_creation_many_keys[50] 0.1809ms 89.6258μs 11.1575 KOps/s 11.0484 KOps/s $\color{#35bf28}+0.99\%$
test_creation_many_keys[100] 0.2240ms 0.1785ms 5.6032 KOps/s 5.6130 KOps/s $\color{#d91a1a}-0.17\%$
test_creation_nested_many_keys[10] 75.0530μs 45.3626μs 22.0446 KOps/s 21.9484 KOps/s $\color{#35bf28}+0.44\%$
test_creation_nested_many_keys[50] 0.2485ms 0.1871ms 5.3454 KOps/s 5.4287 KOps/s $\color{#d91a1a}-1.53\%$
test_clone 40.4010μs 13.4057μs 74.5954 KOps/s 73.1322 KOps/s $\color{#35bf28}+2.00\%$
test_getitem[int] 1.5501ms 15.2959μs 65.3771 KOps/s 65.4130 KOps/s $\color{#d91a1a}-0.06\%$
test_getitem[slice_int] 0.1401ms 24.6924μs 40.4982 KOps/s 39.9348 KOps/s $\color{#35bf28}+1.41\%$
test_getitem[range] 0.1762ms 65.1469μs 15.3499 KOps/s 15.1992 KOps/s $\color{#35bf28}+0.99\%$
test_getitem[tuple] 0.1609ms 24.1741μs 41.3665 KOps/s 41.0801 KOps/s $\color{#35bf28}+0.70\%$
test_getitem[list] 0.1811ms 59.7976μs 16.7231 KOps/s 16.5696 KOps/s $\color{#35bf28}+0.93\%$
test_setitem_dim[int] 48.2610μs 26.4896μs 37.7507 KOps/s 38.4427 KOps/s $\color{#d91a1a}-1.80\%$
test_setitem_dim[slice_int] 86.5820μs 45.7477μs 21.8590 KOps/s 23.2835 KOps/s $\textbf{\color{#d91a1a}-6.12\%}$
test_setitem_dim[range] 0.1509ms 99.8150μs 10.0185 KOps/s 10.4371 KOps/s $\color{#d91a1a}-4.01\%$
test_setitem_dim[tuple] 66.1320μs 41.6081μs 24.0338 KOps/s 23.9569 KOps/s $\color{#35bf28}+0.32\%$
test_setitem 51.0510μs 17.6330μs 56.7120 KOps/s 54.6742 KOps/s $\color{#35bf28}+3.73\%$
test_set 59.7910μs 17.0124μs 58.7806 KOps/s 57.1974 KOps/s $\color{#35bf28}+2.77\%$
test_set_shared 0.5081ms 0.2169ms 4.6104 KOps/s 4.7568 KOps/s $\color{#d91a1a}-3.08\%$
test_update 0.3477ms 21.9257μs 45.6085 KOps/s 44.9352 KOps/s $\color{#35bf28}+1.50\%$
test_update_nested 78.0830μs 33.0091μs 30.2947 KOps/s 29.4657 KOps/s $\color{#35bf28}+2.81\%$
test_update__nested 0.5000ms 33.8615μs 29.5321 KOps/s 29.2477 KOps/s $\color{#35bf28}+0.97\%$
test_set_nested 46.4610μs 18.8633μs 53.0131 KOps/s 51.5311 KOps/s $\color{#35bf28}+2.88\%$
test_set_nested_new 0.2243ms 23.4944μs 42.5634 KOps/s 40.6019 KOps/s $\color{#35bf28}+4.83\%$
test_select 83.0320μs 40.3478μs 24.7845 KOps/s 23.7819 KOps/s $\color{#35bf28}+4.22\%$
test_select_nested 0.1235ms 74.5629μs 13.4115 KOps/s 13.3960 KOps/s $\color{#35bf28}+0.12\%$
test_exclude_nested 0.1526ms 91.3696μs 10.9446 KOps/s 10.9107 KOps/s $\color{#35bf28}+0.31\%$
test_empty[True] 0.4375ms 0.3999ms 2.5008 KOps/s 2.4881 KOps/s $\color{#35bf28}+0.51\%$
test_empty[False] 7.8128μs 1.3715μs 729.1030 KOps/s 706.8595 KOps/s $\color{#35bf28}+3.15\%$
test_to 96.2920μs 69.4445μs 14.4000 KOps/s 14.2442 KOps/s $\color{#35bf28}+1.09\%$
test_to_nonblocking 0.2039ms 71.1318μs 14.0584 KOps/s 14.8152 KOps/s $\textbf{\color{#d91a1a}-5.11\%}$
test_unbind_speed 0.3857ms 0.3350ms 2.9855 KOps/s 2.9972 KOps/s $\color{#d91a1a}-0.39\%$
test_unbind_speed_stack0 0.4150ms 0.3279ms 3.0500 KOps/s 3.0101 KOps/s $\color{#35bf28}+1.33\%$
test_unbind_speed_stack1 0.1201s 0.8489ms 1.1780 KOps/s 1.1624 KOps/s $\color{#35bf28}+1.34\%$
test_split 0.1170s 1.2792ms 781.7264 Ops/s 703.9326 Ops/s $\textbf{\color{#35bf28}+11.05\%}$
test_chunk 1.1477ms 1.0828ms 923.5491 Ops/s 914.0449 Ops/s $\color{#35bf28}+1.04\%$
test_to_cpu_blocking 0.1363s 21.8746ms 45.7152 Ops/s 52.0693 Ops/s $\textbf{\color{#d91a1a}-12.20\%}$
test_to_cpu_global_sync 12.1827ms 11.8472ms 84.4082 Ops/s 85.2317 Ops/s $\color{#d91a1a}-0.97\%$
test_to_cpu_event_sync 13.1911ms 12.7392ms 78.4979 Ops/s 79.5148 Ops/s $\color{#d91a1a}-1.28\%$
test_to_cpu_default 13.1957ms 12.8112ms 78.0567 Ops/s 70.7751 Ops/s $\textbf{\color{#35bf28}+10.29\%}$
test_consolidate[False-None] 4.3941ms 4.2419ms 235.7441 Ops/s 237.3480 Ops/s $\color{#d91a1a}-0.68\%$
test_consolidate[default-None] 3.4824ms 2.1574ms 463.5209 Ops/s 457.9925 Ops/s $\color{#35bf28}+1.21\%$
test_consolidate[reduce-overhead-None] 3.5108ms 2.0726ms 482.4796 Ops/s 472.6723 Ops/s $\color{#35bf28}+2.07\%$
test_consolidate_njt[False-None] 9.8611ms 8.6244ms 115.9497 Ops/s 116.8388 Ops/s $\color{#d91a1a}-0.76\%$
test_to[False-False-None] 2.3145ms 2.1677ms 461.3258 Ops/s 461.9083 Ops/s $\color{#d91a1a}-0.13\%$
test_to[True-False-None] 2.2074ms 1.9069ms 524.4174 Ops/s 530.6744 Ops/s $\color{#d91a1a}-1.18\%$
test_to[within-False-None] 6.5767ms 6.2520ms 159.9492 Ops/s 160.8743 Ops/s $\color{#d91a1a}-0.58\%$
test_to[True-default-None] 10.8686ms 10.6455ms 93.9361 Ops/s 92.3930 Ops/s $\color{#35bf28}+1.67\%$
test_to_njt[False-False-None] 8.7714ms 8.6305ms 115.8675 Ops/s 114.8375 Ops/s $\color{#35bf28}+0.90\%$
test_to_njt[True-False-None] 7.2135ms 6.9638ms 143.6003 Ops/s 144.3230 Ops/s $\color{#d91a1a}-0.50\%$
test_to_njt[within-False-None] 16.7821ms 15.8813ms 62.9671 Ops/s 63.7387 Ops/s $\color{#d91a1a}-1.21\%$
test_creation[device0] 0.4115ms 0.1225ms 8.1610 KOps/s 8.2754 KOps/s $\color{#d91a1a}-1.38\%$
test_creation_from_tensor 0.5201ms 0.1179ms 8.4817 KOps/s 8.4218 KOps/s $\color{#35bf28}+0.71\%$
test_add_one[memmap_tensor0] 0.3819ms 6.8633μs 145.7023 KOps/s 146.3552 KOps/s $\color{#d91a1a}-0.45\%$
test_contiguous[memmap_tensor0] 23.9500μs 0.6617μs 1.5112 MOps/s 2.2902 MOps/s $\textbf{\color{#d91a1a}-34.02\%}$
test_stack[memmap_tensor0] 30.7410μs 4.2345μs 236.1533 KOps/s 237.0736 KOps/s $\color{#d91a1a}-0.39\%$
test_memmaptd_index 1.2935ms 0.2744ms 3.6446 KOps/s 3.6700 KOps/s $\color{#d91a1a}-0.69\%$
test_memmaptd_index_astensor 0.5408ms 0.3830ms 2.6110 KOps/s 2.6406 KOps/s $\color{#d91a1a}-1.12\%$
test_memmaptd_index_op 0.8970ms 0.6452ms 1.5498 KOps/s 1.5405 KOps/s $\color{#35bf28}+0.60\%$
test_serialize_model 0.1374s 0.1354s 7.3846 Ops/s 7.3205 Ops/s $\color{#35bf28}+0.88\%$
test_serialize_model_pickle 1.3473s 1.2102s 0.8263 Ops/s 0.8266 Ops/s $\color{#d91a1a}-0.04\%$
test_serialize_weights 0.1380s 0.1355s 7.3793 Ops/s 7.4725 Ops/s $\color{#d91a1a}-1.25\%$
test_serialize_weights_returnearly 66.1378ms 62.0958ms 16.1042 Ops/s 11.8683 Ops/s $\textbf{\color{#35bf28}+35.69\%}$
test_serialize_weights_pickle 1.3571s 1.2118s 0.8252 Ops/s 0.8228 Ops/s $\color{#35bf28}+0.29\%$
test_reshape_pytree 0.2267ms 32.7837μs 30.5030 KOps/s 30.0363 KOps/s $\color{#35bf28}+1.55\%$
test_reshape_td 82.4220μs 45.9492μs 21.7632 KOps/s 22.2776 KOps/s $\color{#d91a1a}-2.31\%$
test_view_pytree 0.2135ms 31.8568μs 31.3905 KOps/s 30.8893 KOps/s $\color{#35bf28}+1.62\%$
test_view_td 94.3530μs 53.8833μs 18.5586 KOps/s 18.4089 KOps/s $\color{#35bf28}+0.81\%$
test_unbind_pytree 0.2341ms 36.5024μs 27.3954 KOps/s 27.1380 KOps/s $\color{#35bf28}+0.95\%$
test_unbind_td 0.1155ms 50.5690μs 19.7750 KOps/s 19.9674 KOps/s $\color{#d91a1a}-0.96\%$
test_split_pytree 0.2445ms 42.2784μs 23.6527 KOps/s 23.2188 KOps/s $\color{#35bf28}+1.87\%$
test_split_td 0.2326ms 64.5738μs 15.4861 KOps/s 15.5685 KOps/s $\color{#d91a1a}-0.53\%$
test_add_pytree 0.2292ms 41.9338μs 23.8471 KOps/s 23.2310 KOps/s $\color{#35bf28}+2.65\%$
test_add_td 96.4130μs 55.7356μs 17.9418 KOps/s 18.0173 KOps/s $\color{#d91a1a}-0.42\%$
test_compile_add_one_nested[tensordict-compile] 0.4228ms 0.2239ms 4.4660 KOps/s 4.4208 KOps/s $\color{#35bf28}+1.02\%$
test_compile_add_one_nested[tensordict-eager] 0.3087ms 0.2021ms 4.9481 KOps/s 4.9908 KOps/s $\color{#d91a1a}-0.85\%$
test_compile_add_one_nested[pytree-compile] 0.3200ms 0.1344ms 7.4414 KOps/s 7.0189 KOps/s $\textbf{\color{#35bf28}+6.02\%}$
test_compile_add_one_nested[pytree-eager] 0.4587ms 0.1750ms 5.7127 KOps/s 5.6050 KOps/s $\color{#35bf28}+1.92\%$
test_compile_copy_nested[tensordict-compile] 72.7420μs 15.4546μs 64.7057 KOps/s 65.3725 KOps/s $\color{#d91a1a}-1.02\%$
test_compile_copy_nested[tensordict-eager] 0.2987ms 53.3399μs 18.7477 KOps/s 18.4925 KOps/s $\color{#35bf28}+1.38\%$
test_compile_copy_nested[pytree-compile] 0.1783ms 15.6276μs 63.9893 KOps/s 63.7866 KOps/s $\color{#35bf28}+0.32\%$
test_compile_copy_nested[pytree-eager] 0.3680ms 67.5062μs 14.8135 KOps/s 14.6668 KOps/s $\color{#35bf28}+1.00\%$
test_compile_add_one_flat[tensordict-compile] 0.4103ms 0.2479ms 4.0341 KOps/s 3.6627 KOps/s $\textbf{\color{#35bf28}+10.14\%}$
test_compile_add_one_flat[tensordict-eager] 0.3400ms 0.2805ms 3.5655 KOps/s 3.4770 KOps/s $\color{#35bf28}+2.54\%$
test_compile_add_one_flat[tensorclass-compile] 0.2511ms 0.1862ms 5.3720 KOps/s 5.1378 KOps/s $\color{#35bf28}+4.56\%$
test_compile_add_one_flat[tensorclass-eager] 0.2189ms 74.6105μs 13.4029 KOps/s 13.4333 KOps/s $\color{#d91a1a}-0.23\%$
test_compile_add_one_flat[pytree-compile] 0.2859ms 0.1922ms 5.2017 KOps/s 4.9978 KOps/s $\color{#35bf28}+4.08\%$
test_compile_add_one_flat[pytree-eager] 0.7996ms 0.5147ms 1.9430 KOps/s 1.9462 KOps/s $\color{#d91a1a}-0.16\%$
test_compile_add_self_flat[tensordict-eager] 0.4029ms 0.3277ms 3.0514 KOps/s 3.0108 KOps/s $\color{#35bf28}+1.35\%$
test_compile_add_self_flat[tensordict-compile] 0.3865ms 0.2525ms 3.9600 KOps/s 3.7829 KOps/s $\color{#35bf28}+4.68\%$
test_compile_add_self_flat[tensorclass-eager] 0.2125ms 91.0658μs 10.9811 KOps/s 11.1368 KOps/s $\color{#d91a1a}-1.40\%$
test_compile_add_self_flat[tensorclass-compile] 0.3896ms 0.1881ms 5.3173 KOps/s 5.1347 KOps/s $\color{#35bf28}+3.56\%$
test_compile_add_self_flat[pytree-eager] 0.6344ms 0.4202ms 2.3795 KOps/s 2.3919 KOps/s $\color{#d91a1a}-0.52\%$
test_compile_add_self_flat[pytree-compile] 0.5058ms 0.2535ms 3.9454 KOps/s 5.0254 KOps/s $\textbf{\color{#d91a1a}-21.49\%}$
test_compile_copy_flat[tensordict-compile] 96.8720μs 19.6424μs 50.9103 KOps/s 50.8265 KOps/s $\color{#35bf28}+0.16\%$
test_compile_copy_flat[tensordict-eager] 81.0930μs 41.4483μs 24.1264 KOps/s 24.5853 KOps/s $\color{#d91a1a}-1.87\%$
test_compile_copy_flat[pytree-compile] 0.1891ms 21.0838μs 47.4298 KOps/s 46.0667 KOps/s $\color{#35bf28}+2.96\%$
test_compile_copy_flat[pytree-eager] 0.3274ms 52.3136μs 19.1155 KOps/s 18.9577 KOps/s $\color{#35bf28}+0.83\%$
test_compile_assign_and_add[tensordict-compile] 2.2112ms 0.2403ms 4.1616 KOps/s 3.9710 KOps/s $\color{#35bf28}+4.80\%$
test_compile_assign_and_add[tensordict-eager] 3.4507ms 3.2434ms 308.3167 Ops/s 306.7844 Ops/s $\color{#35bf28}+0.50\%$
test_compile_assign_and_add[pytree-compile] 2.1180ms 0.1801ms 5.5521 KOps/s 5.4473 KOps/s $\color{#35bf28}+1.92\%$
test_compile_assign_and_add[pytree-eager] 2.9972ms 2.7421ms 364.6849 Ops/s 359.1622 Ops/s $\color{#35bf28}+1.54\%$
test_compile_indexing[tensor-tensordict-compile] 0.2977ms 0.1345ms 7.4357 KOps/s 7.3047 KOps/s $\color{#35bf28}+1.79\%$
test_compile_indexing[tensor-tensordict-eager] 0.3176ms 74.4679μs 13.4286 KOps/s 12.8182 KOps/s $\color{#35bf28}+4.76\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2421ms 0.1197ms 8.3557 KOps/s 8.3219 KOps/s $\color{#35bf28}+0.41\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2634ms 44.3852μs 22.5300 KOps/s 22.3574 KOps/s $\color{#35bf28}+0.77\%$
test_compile_indexing[tensor-pytree-compile] 0.1777ms 0.1200ms 8.3338 KOps/s 8.3533 KOps/s $\color{#d91a1a}-0.23\%$
test_compile_indexing[tensor-pytree-eager] 0.2655ms 44.4820μs 22.4810 KOps/s 22.3806 KOps/s $\color{#35bf28}+0.45\%$
test_compile_indexing[slice-tensordict-compile] 0.2085ms 82.6429μs 12.1002 KOps/s 12.0173 KOps/s $\color{#35bf28}+0.69\%$
test_compile_indexing[slice-tensordict-eager] 0.2067ms 27.8925μs 35.8520 KOps/s 35.8097 KOps/s $\color{#35bf28}+0.12\%$
test_compile_indexing[slice-tensorclass-compile] 0.1772ms 70.0984μs 14.2657 KOps/s 14.2859 KOps/s $\color{#d91a1a}-0.14\%$
test_compile_indexing[slice-tensorclass-eager] 0.2354ms 21.7806μs 45.9124 KOps/s 44.2131 KOps/s $\color{#35bf28}+3.84\%$
test_compile_indexing[slice-pytree-compile] 0.1450ms 70.2292μs 14.2391 KOps/s 14.0389 KOps/s $\color{#35bf28}+1.43\%$
test_compile_indexing[slice-pytree-eager] 0.2901ms 21.7488μs 45.9795 KOps/s 44.5756 KOps/s $\color{#35bf28}+3.15\%$
test_compile_indexing[int-tensordict-compile] 0.1310ms 83.9917μs 11.9059 KOps/s 11.9874 KOps/s $\color{#d91a1a}-0.68\%$
test_compile_indexing[int-tensordict-eager] 0.2223ms 27.7825μs 35.9939 KOps/s 36.5047 KOps/s $\color{#d91a1a}-1.40\%$
test_compile_indexing[int-tensorclass-compile] 0.1221ms 70.8558μs 14.1132 KOps/s 14.1058 KOps/s $\color{#35bf28}+0.05\%$
test_compile_indexing[int-tensorclass-eager] 0.2404ms 22.0443μs 45.3631 KOps/s 44.5404 KOps/s $\color{#35bf28}+1.85\%$
test_compile_indexing[int-pytree-compile] 0.1167ms 70.8671μs 14.1109 KOps/s 13.9458 KOps/s $\color{#35bf28}+1.18\%$
test_compile_indexing[int-pytree-eager] 0.2784ms 21.8399μs 45.7877 KOps/s 44.8863 KOps/s $\color{#35bf28}+2.01\%$
test_compile_replace[single-eager] 91.5020μs 48.6412μs 20.5587 KOps/s 20.4968 KOps/s $\color{#35bf28}+0.30\%$
test_compile_replace[single-compile] 0.2221ms 0.1309ms 7.6385 KOps/s 7.7239 KOps/s $\color{#d91a1a}-1.11\%$
test_compile_replace[multi-eager] 0.6228ms 0.5621ms 1.7790 KOps/s 1.7894 KOps/s $\color{#d91a1a}-0.58\%$
test_compile_replace[multi-compile] 0.2616ms 0.1382ms 7.2368 KOps/s 7.0503 KOps/s $\color{#35bf28}+2.64\%$
test_compile_tc_getattr_20[eager] 0.2079ms 0.1629ms 6.1397 KOps/s 6.3352 KOps/s $\color{#d91a1a}-3.09\%$
test_compile_tc_getattr_20[compile] 0.3736ms 0.1483ms 6.7428 KOps/s 6.4028 KOps/s $\textbf{\color{#35bf28}+5.31\%}$
test_compile_clone_shallow[20-eager] 0.1073ms 19.1564μs 52.2018 KOps/s 51.7685 KOps/s $\color{#35bf28}+0.84\%$
test_compile_clone_shallow[20-compile] 80.1630μs 16.7264μs 59.7858 KOps/s 59.2972 KOps/s $\color{#35bf28}+0.82\%$
test_compile_clone_shallow[40-eager] 80.5220μs 33.5650μs 29.7929 KOps/s 29.5133 KOps/s $\color{#35bf28}+0.95\%$
test_compile_clone_shallow[40-compile] 54.5310μs 18.1533μs 55.0865 KOps/s 52.1218 KOps/s $\textbf{\color{#35bf28}+5.69\%}$
test_compile_clone_shallow[80-eager] 98.2230μs 62.2821μs 16.0560 KOps/s 15.9887 KOps/s $\color{#35bf28}+0.42\%$
test_compile_clone_shallow[80-compile] 88.6320μs 21.4401μs 46.6416 KOps/s 45.5289 KOps/s $\color{#35bf28}+2.44\%$
test_compile_update_inplace[eager] 0.1461ms 59.8405μs 16.7111 KOps/s 16.3127 KOps/s $\color{#35bf28}+2.44\%$
test_compile_update_inplace[compile] 0.4346ms 0.1465ms 6.8258 KOps/s 6.6872 KOps/s $\color{#35bf28}+2.07\%$
test_mod_add[eager] 0.3110ms 0.2273ms 4.4000 KOps/s 4.4188 KOps/s $\color{#d91a1a}-0.43\%$
test_mod_add[compile] 0.3966ms 0.1304ms 7.6704 KOps/s 7.4600 KOps/s $\color{#35bf28}+2.82\%$
test_mod_add[compile-overhead] 0.2626ms 0.1773ms 5.6409 KOps/s 5.5433 KOps/s $\color{#35bf28}+1.76\%$
test_mod_wrap[eager] 0.5462ms 0.4677ms 2.1379 KOps/s 2.1188 KOps/s $\color{#35bf28}+0.90\%$
test_mod_wrap[compile] 0.8761ms 0.4024ms 2.4852 KOps/s 2.5095 KOps/s $\color{#d91a1a}-0.97\%$
test_mod_wrap[compile-overhead] 8.6029ms 4.7337ms 211.2529 Ops/s 215.7297 Ops/s $\color{#d91a1a}-2.08\%$
test_mod_wrap_and_backward[eager] 2.1107ms 1.6919ms 591.0437 Ops/s 589.0452 Ops/s $\color{#35bf28}+0.34\%$
test_mod_wrap_and_backward[compile] 2.4392ms 1.6101ms 621.0947 Ops/s 644.0919 Ops/s $\color{#d91a1a}-3.57\%$
test_mod_wrap_and_backward[compile-overhead] 1.3578ms 0.9508ms 1.0517 KOps/s 1.0397 KOps/s $\color{#35bf28}+1.16\%$
test_seq_add[eager] 0.9886ms 0.5158ms 1.9389 KOps/s 1.9632 KOps/s $\color{#d91a1a}-1.24\%$
test_seq_add[compile] 0.5676ms 0.1426ms 7.0108 KOps/s 6.8228 KOps/s $\color{#35bf28}+2.76\%$
test_seq_add[compile-overhead] 0.3097ms 0.1886ms 5.3030 KOps/s 5.2420 KOps/s $\color{#35bf28}+1.16\%$
test_seq_wrap[eager] 1.9034ms 1.4300ms 699.3200 Ops/s 688.8125 Ops/s $\color{#35bf28}+1.53\%$
test_seq_wrap[compile] 0.5128ms 0.4267ms 2.3438 KOps/s 2.3103 KOps/s $\color{#35bf28}+1.45\%$
test_seq_wrap[compile-overhead] 0.4019ms 0.3102ms 3.2233 KOps/s 3.1420 KOps/s $\color{#35bf28}+2.59\%$
test_func_call_runtime[False-eager] 0.9679ms 0.8382ms 1.1930 KOps/s 1.1824 KOps/s $\color{#35bf28}+0.90\%$
test_func_call_runtime[False-compile] 1.1466ms 0.9996ms 1.0004 KOps/s 1.0054 KOps/s $\color{#d91a1a}-0.50\%$
test_func_call_runtime[False-compile-overhead] 0.6326ms 0.5136ms 1.9469 KOps/s 1.9366 KOps/s $\color{#35bf28}+0.53\%$
test_func_call_runtime[True-eager] 1.2249ms 1.0701ms 934.4599 Ops/s 940.1313 Ops/s $\color{#d91a1a}-0.60\%$
test_func_call_runtime[True-compile] 1.5169ms 1.0101ms 990.0378 Ops/s 1.0100 KOps/s $\color{#d91a1a}-1.97\%$
test_func_call_runtime[True-compile-overhead] 0.9748ms 0.5275ms 1.8958 KOps/s 1.8709 KOps/s $\color{#35bf28}+1.33\%$
test_func_call_cm_runtime[False-eager] 1.3224ms 0.8364ms 1.1955 KOps/s 1.2513 KOps/s $\color{#d91a1a}-4.46\%$
test_func_call_cm_runtime[False-compile] 1.2011ms 0.9968ms 1.0032 KOps/s 1.0379 KOps/s $\color{#d91a1a}-3.34\%$
test_func_call_cm_runtime[False-compile-overhead] 0.6679ms 0.4910ms 2.0367 KOps/s 2.0191 KOps/s $\color{#35bf28}+0.87\%$
test_func_call_cm_runtime[True-eager] 1.3804ms 1.2175ms 821.3248 Ops/s 823.5949 Ops/s $\color{#d91a1a}-0.28\%$
test_func_call_cm_runtime[True-compile] 1.5216ms 1.0513ms 951.2340 Ops/s 980.4341 Ops/s $\color{#d91a1a}-2.98\%$
test_func_call_cm_runtime[True-compile-overhead] 0.9941ms 0.5490ms 1.8216 KOps/s 1.7976 KOps/s $\color{#35bf28}+1.34\%$
test_vmap_func_call_cm_runtime[eager] 2.9007ms 2.4228ms 412.7379 Ops/s 414.1035 Ops/s $\color{#d91a1a}-0.33\%$
test_vmap_func_call_cm_runtime[compile] 3.3167ms 2.6406ms 378.6999 Ops/s 384.1120 Ops/s $\color{#d91a1a}-1.41\%$
test_vmap_func_call_cm_runtime[compile-overhead] 3.5315ms 2.6085ms 383.3569 Ops/s 384.1900 Ops/s $\color{#d91a1a}-0.22\%$
test_distributed 0.5758ms 0.1483ms 6.7412 KOps/s 6.5727 KOps/s $\color{#35bf28}+2.56\%$
test_tdmodule 0.4962ms 0.1945ms 5.1401 KOps/s 5.0881 KOps/s $\color{#35bf28}+1.02\%$
test_tdmodule_dispatch 0.4532ms 0.3753ms 2.6644 KOps/s 2.6545 KOps/s $\color{#35bf28}+0.37\%$
test_tdseq 46.6210μs 26.6738μs 37.4899 KOps/s 36.9522 KOps/s $\color{#35bf28}+1.46\%$
test_tdseq_dispatch 0.2413ms 0.2170ms 4.6083 KOps/s 4.6214 KOps/s $\color{#d91a1a}-0.28\%$
test_instantiation_functorch 2.1863ms 2.0793ms 480.9291 Ops/s 481.9543 Ops/s $\color{#d91a1a}-0.21\%$
test_exec_functorch 0.2142ms 0.1843ms 5.4268 KOps/s 5.4372 KOps/s $\color{#d91a1a}-0.19\%$
test_exec_functional_call 0.2605ms 0.1639ms 6.1011 KOps/s 6.1337 KOps/s $\color{#d91a1a}-0.53\%$
test_exec_td_decorator 0.5017ms 0.2520ms 3.9689 KOps/s 3.9657 KOps/s $\color{#35bf28}+0.08\%$
test_vmap_mlp_speed_decorator[True-True] 1.2291ms 1.0403ms 961.2840 Ops/s 956.3497 Ops/s $\color{#35bf28}+0.52\%$
test_vmap_mlp_speed_decorator[True-False] 1.2232ms 1.0379ms 963.4618 Ops/s 960.5247 Ops/s $\color{#35bf28}+0.31\%$
test_vmap_mlp_speed_decorator[False-True] 0.9219ms 0.7366ms 1.3576 KOps/s 1.3585 KOps/s $\color{#d91a1a}-0.06\%$
test_vmap_mlp_speed_decorator[False-False] 0.9408ms 0.7394ms 1.3524 KOps/s 1.3604 KOps/s $\color{#d91a1a}-0.59\%$
test_vmap_transformer_speed_decorator[True-True] 21.2816ms 21.1155ms 47.3586 Ops/s 47.2019 Ops/s $\color{#35bf28}+0.33\%$
test_vmap_transformer_speed_decorator[True-False] 21.9580ms 21.3002ms 46.9480 Ops/s 47.2682 Ops/s $\color{#d91a1a}-0.68\%$
test_vmap_transformer_speed_decorator[False-True] 21.6522ms 20.7780ms 48.1279 Ops/s 48.0246 Ops/s $\color{#35bf28}+0.22\%$
test_vmap_transformer_speed_decorator[False-False] 21.6600ms 20.7646ms 48.1590 Ops/s 48.0740 Ops/s $\color{#35bf28}+0.18\%$
test_to_module_speed[True] 2.2381ms 1.6449ms 607.9264 Ops/s 603.3644 Ops/s $\color{#35bf28}+0.76\%$
test_to_module_speed[False] 2.2043ms 1.6037ms 623.5660 Ops/s 611.4969 Ops/s $\color{#35bf28}+1.97\%$
test_to_module_plain_tensor_speed[None] 2.2844ms 1.7361ms 575.9994 Ops/s 569.0413 Ops/s $\color{#35bf28}+1.22\%$
test_to_module_plain_tensor_speed[False] 2.0875ms 1.5769ms 634.1358 Ops/s 623.5956 Ops/s $\color{#35bf28}+1.69\%$
test_to_module_plain_tensor_speed[True] 3.5236ms 3.4064ms 293.5657 Ops/s 289.4474 Ops/s $\color{#35bf28}+1.42\%$
test_tc_init 75.0420μs 43.1847μs 23.1563 KOps/s 23.2469 KOps/s $\color{#d91a1a}-0.39\%$
test_tc_init_tensor_only 28.6500μs 8.9820μs 111.3341 KOps/s 109.7515 KOps/s $\color{#35bf28}+1.44\%$
test_tc_init_nested 0.1183ms 83.2594μs 12.0107 KOps/s 11.6107 KOps/s $\color{#35bf28}+3.44\%$
test_tc_init_many_fields 36.3410μs 14.7495μs 67.7989 KOps/s 65.5487 KOps/s $\color{#35bf28}+3.43\%$
test_tc_first_layer_tensor 23.7210μs 1.8155μs 550.8091 KOps/s 547.3294 KOps/s $\color{#35bf28}+0.64\%$
test_tc_first_layer_tensor_only 1.6010μs 0.4018μs 2.4888 MOps/s 2.5169 MOps/s $\color{#d91a1a}-1.12\%$
test_tc_first_layer_tensor_set 40.3210μs 3.9644μs 252.2467 KOps/s 251.7997 KOps/s $\color{#35bf28}+0.18\%$
test_tc_first_layer_tensor_only_set 43.1410μs 3.3054μs 302.5364 KOps/s 301.1364 KOps/s $\color{#35bf28}+0.46\%$
test_tc_first_layer_nontensor 26.2910μs 6.2494μs 160.0144 KOps/s 155.7199 KOps/s $\color{#35bf28}+2.76\%$
test_tc_second_layer_tensor 32.4310μs 4.4193μs 226.2780 KOps/s 226.1031 KOps/s $\color{#35bf28}+0.08\%$
test_tc_second_layer_nontensor 39.0210μs 8.7094μs 114.8186 KOps/s 110.3966 KOps/s $\color{#35bf28}+4.01\%$
test_unbind 0.2805s 13.8727ms 72.0839 Ops/s 68.7081 Ops/s $\color{#35bf28}+4.91\%$
test_full_like 4.7900ms 4.4933ms 222.5531 Ops/s 223.1230 Ops/s $\color{#d91a1a}-0.26\%$
test_zeros_like 4.8260ms 4.4632ms 224.0561 Ops/s 225.3790 Ops/s $\color{#d91a1a}-0.59\%$
test_ones_like 5.0668ms 4.5102ms 221.7189 Ops/s 223.7804 Ops/s $\color{#d91a1a}-0.92\%$
test_clone 7.6353ms 6.9047ms 144.8297 Ops/s 145.9274 Ops/s $\color{#d91a1a}-0.75\%$
test_squeeze 0.1613ms 14.2419μs 70.2152 KOps/s 68.5645 KOps/s $\color{#35bf28}+2.41\%$
test_unsqueeze 0.1611ms 0.1117ms 8.9564 KOps/s 9.0363 KOps/s $\color{#d91a1a}-0.88\%$
test_split 0.2335ms 0.1812ms 5.5198 KOps/s 5.4976 KOps/s $\color{#35bf28}+0.40\%$
test_permute 0.2605ms 0.2112ms 4.7339 KOps/s 4.8431 KOps/s $\color{#d91a1a}-2.26\%$
test_stack 20.1713ms 19.4334ms 51.4578 Ops/s 22.8595 Ops/s $\textbf{\color{#35bf28}+125.10\%}$
test_cat 20.3155ms 19.4762ms 51.3448 Ops/s 22.7157 Ops/s $\textbf{\color{#35bf28}+126.03\%}$
test_sequential_tensordict 0.6560ms 0.5772ms 1.7325 KOps/s 1.7337 KOps/s $\color{#d91a1a}-0.06\%$
test_sequential_graph_module 0.2724ms 0.1221ms 8.1933 KOps/s 8.2901 KOps/s $\color{#d91a1a}-1.17\%$
test_nested_tensordict 0.9946ms 0.8205ms 1.2188 KOps/s 1.2057 KOps/s $\color{#35bf28}+1.08\%$
test_nested_graph_module 0.2218ms 0.1278ms 7.8239 KOps/s 7.6166 KOps/s $\color{#35bf28}+2.72\%$

@github-actions

github-actions Bot commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

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

Total Benchmarks: 344. Improved: $\large\color{#35bf28}14$. 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.2360ms 0.1646ms 6.0744 KOps/s 5.9392 KOps/s $\color{#35bf28}+2.28\%$
test_collector_preallocated_write[stack_out-cpu-32-8] 0.3894ms 0.2928ms 3.4155 KOps/s 3.4223 KOps/s $\color{#d91a1a}-0.20\%$
test_collector_preallocated_write[stack_out-cpu-32-32] 0.8863ms 0.7738ms 1.2924 KOps/s 1.2784 KOps/s $\color{#35bf28}+1.09\%$
test_collector_preallocated_write[stack_out-cpu-4096-2] 0.3967ms 0.3393ms 2.9475 KOps/s 2.8757 KOps/s $\color{#35bf28}+2.50\%$
test_collector_preallocated_write[stack_out-cpu-4096-8] 1.6883ms 1.1516ms 868.3284 Ops/s 871.0039 Ops/s $\color{#d91a1a}-0.31\%$
test_collector_preallocated_write[stack_out-cpu-4096-32] 8.0226ms 7.6366ms 130.9484 Ops/s 127.2768 Ops/s $\color{#35bf28}+2.88\%$
test_collector_preallocated_write[stack_out-cuda-32-2] 0.3574ms 0.2535ms 3.9447 KOps/s 4.0743 KOps/s $\color{#d91a1a}-3.18\%$
test_collector_preallocated_write[stack_out-cuda-32-8] 0.4921ms 0.3708ms 2.6972 KOps/s 2.6946 KOps/s $\color{#35bf28}+0.09\%$
test_collector_preallocated_write[stack_out-cuda-32-32] 0.9733ms 0.8192ms 1.2208 KOps/s 1.2112 KOps/s $\color{#35bf28}+0.79\%$
test_collector_preallocated_write[stack_out-cuda-4096-2] 0.4063ms 0.2512ms 3.9811 KOps/s 3.8819 KOps/s $\color{#35bf28}+2.56\%$
test_collector_preallocated_write[stack_out-cuda-4096-8] 0.4566ms 0.3743ms 2.6716 KOps/s 2.6937 KOps/s $\color{#d91a1a}-0.82\%$
test_collector_preallocated_write[stack_out-cuda-4096-32] 0.9907ms 0.8850ms 1.1299 KOps/s 1.1308 KOps/s $\color{#d91a1a}-0.07\%$
test_collector_preallocated_write[set_at_loop-cpu-32-2] 0.2687ms 0.2142ms 4.6679 KOps/s 4.6471 KOps/s $\color{#35bf28}+0.45\%$
test_collector_preallocated_write[set_at_loop-cpu-32-8] 0.9260ms 0.8252ms 1.2118 KOps/s 1.1965 KOps/s $\color{#35bf28}+1.28\%$
test_collector_preallocated_write[set_at_loop-cpu-32-32] 3.3567ms 3.2662ms 306.1647 Ops/s 299.6438 Ops/s $\color{#35bf28}+2.18\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-2] 0.5658ms 0.3694ms 2.7073 KOps/s 2.6654 KOps/s $\color{#35bf28}+1.57\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-8] 2.5981ms 1.7634ms 567.1005 Ops/s 590.3840 Ops/s $\color{#d91a1a}-3.94\%$
test_collector_preallocated_write[set_at_loop-cpu-4096-32] 11.0906ms 10.6461ms 93.9307 Ops/s 89.8316 Ops/s $\color{#35bf28}+4.56\%$
test_collector_preallocated_write[set_at_loop-cuda-32-2] 0.5672ms 0.4023ms 2.4858 KOps/s 2.5588 KOps/s $\color{#d91a1a}-2.85\%$
test_collector_preallocated_write[set_at_loop-cuda-32-8] 1.9813ms 1.4688ms 680.8119 Ops/s 689.7900 Ops/s $\color{#d91a1a}-1.30\%$
test_collector_preallocated_write[set_at_loop-cuda-32-32] 6.0558ms 5.7644ms 173.4787 Ops/s 177.9222 Ops/s $\color{#d91a1a}-2.50\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-2] 0.5438ms 0.4014ms 2.4913 KOps/s 2.5678 KOps/s $\color{#d91a1a}-2.98\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-8] 1.7059ms 1.5022ms 665.6836 Ops/s 694.8359 Ops/s $\color{#d91a1a}-4.20\%$
test_collector_preallocated_write[set_at_loop-cuda-4096-32] 6.0818ms 5.8455ms 171.0712 Ops/s 175.8295 Ops/s $\color{#d91a1a}-2.71\%$
test_collector_preallocated_write[copy_at-cpu-32-2] 0.5907ms 0.1730ms 5.7804 KOps/s 5.6694 KOps/s $\color{#35bf28}+1.96\%$
test_collector_preallocated_write[copy_at-cpu-32-8] 0.7963ms 0.6771ms 1.4769 KOps/s 1.4717 KOps/s $\color{#35bf28}+0.35\%$
test_collector_preallocated_write[copy_at-cpu-32-32] 2.8601ms 2.7262ms 366.8070 Ops/s 366.1316 Ops/s $\color{#35bf28}+0.18\%$
test_collector_preallocated_write[copy_at-cpu-4096-2] 0.7485ms 0.3171ms 3.1540 KOps/s 3.1157 KOps/s $\color{#35bf28}+1.23\%$
test_collector_preallocated_write[copy_at-cpu-4096-8] 1.9996ms 1.6372ms 610.7984 Ops/s 659.7256 Ops/s $\textbf{\color{#d91a1a}-7.42\%}$
test_collector_preallocated_write[copy_at-cpu-4096-32] 10.5154ms 9.9088ms 100.9199 Ops/s 99.5342 Ops/s $\color{#35bf28}+1.39\%$
test_collector_preallocated_write[copy_at-cuda-32-2] 0.4757ms 0.3233ms 3.0927 KOps/s 3.0627 KOps/s $\color{#35bf28}+0.98\%$
test_collector_preallocated_write[copy_at-cuda-32-8] 1.3598ms 1.2066ms 828.7437 Ops/s 835.6945 Ops/s $\color{#d91a1a}-0.83\%$
test_collector_preallocated_write[copy_at-cuda-32-32] 5.0786ms 4.7123ms 212.2128 Ops/s 214.2685 Ops/s $\color{#d91a1a}-0.96\%$
test_collector_preallocated_write[copy_at-cuda-4096-2] 0.4768ms 0.3256ms 3.0716 KOps/s 3.0970 KOps/s $\color{#d91a1a}-0.82\%$
test_collector_preallocated_write[copy_at-cuda-4096-8] 1.4456ms 1.2124ms 824.8046 Ops/s 831.2155 Ops/s $\color{#d91a1a}-0.77\%$
test_collector_preallocated_write[copy_at-cuda-4096-32] 4.9067ms 4.7346ms 211.2133 Ops/s 213.2174 Ops/s $\color{#d91a1a}-0.94\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-2] 0.1556ms 0.1201ms 8.3237 KOps/s 8.3821 KOps/s $\color{#d91a1a}-0.70\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-8] 0.5094ms 0.4687ms 2.1338 KOps/s 2.1513 KOps/s $\color{#d91a1a}-0.81\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-32-32] 1.9990ms 1.8715ms 534.3234 Ops/s 541.6635 Ops/s $\color{#d91a1a}-1.36\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-2] 0.5050ms 0.2557ms 3.9111 KOps/s 3.9781 KOps/s $\color{#d91a1a}-1.69\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-8] 1.7201ms 1.2634ms 791.5367 Ops/s 814.6852 Ops/s $\color{#d91a1a}-2.84\%$
test_collector_preallocated_write[manual_leaf_copy-cpu-4096-32] 9.4775ms 8.9781ms 111.3818 Ops/s 114.4222 Ops/s $\color{#d91a1a}-2.66\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-2] 0.4600ms 0.2769ms 3.6110 KOps/s 3.6601 KOps/s $\color{#d91a1a}-1.34\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-8] 1.1886ms 1.0208ms 979.6495 Ops/s 988.8233 Ops/s $\color{#d91a1a}-0.93\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-32-32] 4.1917ms 3.9841ms 250.9973 Ops/s 252.8920 Ops/s $\color{#d91a1a}-0.75\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-2] 0.4585ms 0.2747ms 3.6408 KOps/s 3.6710 KOps/s $\color{#d91a1a}-0.82\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-8] 1.3499ms 1.0374ms 963.9272 Ops/s 982.4931 Ops/s $\color{#d91a1a}-1.89\%$
test_collector_preallocated_write[manual_leaf_copy-cuda-4096-32] 4.2117ms 3.9928ms 250.4498 Ops/s 251.7301 Ops/s $\color{#d91a1a}-0.51\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-128] 17.0866ms 16.3605ms 61.1230 Ops/s 60.3773 Ops/s $\color{#35bf28}+1.23\%$
test_wide_nested_preallocated_write[stack_out-cpu-128-256] 33.2625ms 32.8884ms 30.4059 Ops/s 30.2300 Ops/s $\color{#35bf28}+0.58\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-128] 89.8132ms 88.1628ms 11.3427 Ops/s 11.5916 Ops/s $\color{#d91a1a}-2.15\%$
test_wide_nested_preallocated_write[stack_out-cpu-4096-256] 0.2121s 0.2100s 4.7618 Ops/s 4.9304 Ops/s $\color{#d91a1a}-3.42\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-128] 13.9855ms 13.7129ms 72.9239 Ops/s 72.5781 Ops/s $\color{#35bf28}+0.48\%$
test_wide_nested_preallocated_write[stack_out-cuda-128-256] 27.3468ms 27.1432ms 36.8416 Ops/s 37.2725 Ops/s $\color{#d91a1a}-1.16\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-128] 14.0707ms 13.9480ms 71.6948 Ops/s 71.9212 Ops/s $\color{#d91a1a}-0.31\%$
test_wide_nested_preallocated_write[stack_out-cuda-4096-256] 27.6642ms 27.5008ms 36.3625 Ops/s 36.3770 Ops/s $\color{#d91a1a}-0.04\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-128] 54.5312ms 52.2575ms 19.1360 Ops/s 19.2568 Ops/s $\color{#d91a1a}-0.63\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-128-256] 0.1047s 0.1041s 9.6049 Ops/s 9.5849 Ops/s $\color{#35bf28}+0.21\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-128] 0.1538s 0.1513s 6.6084 Ops/s 6.8710 Ops/s $\color{#d91a1a}-3.82\%$
test_wide_nested_preallocated_write[set_at_loop-cpu-4096-256] 0.3324s 0.3306s 3.0252 Ops/s 3.0562 Ops/s $\color{#d91a1a}-1.01\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-128] 73.8795ms 73.4017ms 13.6237 Ops/s 13.5965 Ops/s $\color{#35bf28}+0.20\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-128-256] 0.1572s 0.1557s 6.4244 Ops/s 6.8062 Ops/s $\textbf{\color{#d91a1a}-5.61\%}$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-128] 78.9272ms 76.8049ms 13.0200 Ops/s 13.3896 Ops/s $\color{#d91a1a}-2.76\%$
test_wide_nested_preallocated_write[set_at_loop-cuda-4096-256] 0.1480s 0.1473s 6.7897 Ops/s 6.7089 Ops/s $\color{#35bf28}+1.20\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-128] 30.6081ms 30.4138ms 32.8798 Ops/s 32.8224 Ops/s $\color{#35bf28}+0.17\%$
test_wide_nested_preallocated_write[copy_at-cpu-128-256] 61.5511ms 61.2867ms 16.3168 Ops/s 16.1601 Ops/s $\color{#35bf28}+0.97\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-128] 0.1273s 0.1243s 8.0482 Ops/s 8.2314 Ops/s $\color{#d91a1a}-2.23\%$
test_wide_nested_preallocated_write[copy_at-cpu-4096-256] 0.2839s 0.2827s 3.5375 Ops/s 3.7299 Ops/s $\textbf{\color{#d91a1a}-5.16\%}$
test_wide_nested_preallocated_write[copy_at-cuda-128-128] 43.4974ms 43.0432ms 23.2324 Ops/s 23.0198 Ops/s $\color{#35bf28}+0.92\%$
test_wide_nested_preallocated_write[copy_at-cuda-128-256] 86.5481ms 86.3899ms 11.5754 Ops/s 11.5399 Ops/s $\color{#35bf28}+0.31\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-128] 43.6131ms 43.4128ms 23.0347 Ops/s 22.8776 Ops/s $\color{#35bf28}+0.69\%$
test_wide_nested_preallocated_write[copy_at-cuda-4096-256] 87.1305ms 86.4544ms 11.5668 Ops/s 11.5010 Ops/s $\color{#35bf28}+0.57\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-128] 20.3303ms 20.2095ms 49.4816 Ops/s 49.5777 Ops/s $\color{#d91a1a}-0.19\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-128-256] 41.5816ms 41.2726ms 24.2291 Ops/s 24.0962 Ops/s $\color{#35bf28}+0.55\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-128] 0.1113s 0.1106s 9.0388 Ops/s 9.2113 Ops/s $\color{#d91a1a}-1.87\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cpu-4096-256] 0.2602s 0.2526s 3.9581 Ops/s 3.9995 Ops/s $\color{#d91a1a}-1.03\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-128] 36.7035ms 36.3230ms 27.5307 Ops/s 27.7527 Ops/s $\color{#d91a1a}-0.80\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-128-256] 73.3024ms 72.9391ms 13.7101 Ops/s 13.8621 Ops/s $\color{#d91a1a}-1.10\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-128] 36.9471ms 36.1685ms 27.6484 Ops/s 27.6247 Ops/s $\color{#35bf28}+0.09\%$
test_wide_nested_preallocated_write[manual_leaf_copy-cuda-4096-256] 73.7608ms 72.9189ms 13.7139 Ops/s 13.9161 Ops/s $\color{#d91a1a}-1.45\%$
test_plain_set_nested 37.4410μs 15.6551μs 63.8769 KOps/s 65.3141 KOps/s $\color{#d91a1a}-2.20\%$
test_plain_set_stack_nested 36.6710μs 15.3052μs 65.3371 KOps/s 65.7253 KOps/s $\color{#d91a1a}-0.59\%$
test_plain_set_nested_inplace 46.3110μs 17.0466μs 58.6628 KOps/s 58.7766 KOps/s $\color{#d91a1a}-0.19\%$
test_plain_set_stack_nested_inplace 46.5810μs 16.4661μs 60.7309 KOps/s 60.2148 KOps/s $\color{#35bf28}+0.86\%$
test_items 71.7010μs 6.0527μs 165.2167 KOps/s 164.6577 KOps/s $\color{#35bf28}+0.34\%$
test_items_nested 0.5033ms 0.4603ms 2.1723 KOps/s 2.0899 KOps/s $\color{#35bf28}+3.94\%$
test_items_nested_locked 0.5238ms 0.4665ms 2.1436 KOps/s 2.0772 KOps/s $\color{#35bf28}+3.20\%$
test_items_nested_leaf 0.1406ms 97.3394μs 10.2733 KOps/s 10.2266 KOps/s $\color{#35bf28}+0.46\%$
test_items_stack_nested 0.5442ms 0.4584ms 2.1817 KOps/s 2.1188 KOps/s $\color{#35bf28}+2.97\%$
test_items_stack_nested_leaf 0.1731ms 96.4911μs 10.3637 KOps/s 10.2451 KOps/s $\color{#35bf28}+1.16\%$
test_items_stack_nested_locked 0.5235ms 0.4615ms 2.1670 KOps/s 2.0901 KOps/s $\color{#35bf28}+3.68\%$
test_keys 31.1210μs 4.2523μs 235.1678 KOps/s 236.5577 KOps/s $\color{#d91a1a}-0.59\%$
test_keys_nested 0.1770ms 0.1297ms 7.7127 KOps/s 7.5578 KOps/s $\color{#35bf28}+2.05\%$
test_keys_nested_locked 2.1512ms 0.1385ms 7.2223 KOps/s 7.1155 KOps/s $\color{#35bf28}+1.50\%$
test_keys_nested_leaf 0.1638ms 0.1203ms 8.3107 KOps/s 8.1836 KOps/s $\color{#35bf28}+1.55\%$
test_keys_stack_nested 0.1688ms 0.1290ms 7.7548 KOps/s 7.5604 KOps/s $\color{#35bf28}+2.57\%$
test_keys_stack_nested_leaf 0.1651ms 0.1194ms 8.3771 KOps/s 8.2230 KOps/s $\color{#35bf28}+1.87\%$
test_keys_stack_nested_locked 0.1871ms 0.1373ms 7.2838 KOps/s 7.1975 KOps/s $\color{#35bf28}+1.20\%$
test_values 7.8272μs 0.9926μs 1.0074 MOps/s 981.0217 KOps/s $\color{#35bf28}+2.69\%$
test_values_nested 99.2120μs 52.8569μs 18.9190 KOps/s 18.7010 KOps/s $\color{#35bf28}+1.17\%$
test_values_nested_locked 0.2520ms 55.9492μs 17.8733 KOps/s 17.4806 KOps/s $\color{#35bf28}+2.25\%$
test_values_nested_leaf 99.6420μs 60.0506μs 16.6526 KOps/s 16.3262 KOps/s $\color{#35bf28}+2.00\%$
test_values_stack_nested 0.1265ms 52.3699μs 19.0949 KOps/s 18.6805 KOps/s $\color{#35bf28}+2.22\%$
test_values_stack_nested_leaf 92.4510μs 59.8478μs 16.7091 KOps/s 16.3427 KOps/s $\color{#35bf28}+2.24\%$
test_values_stack_nested_locked 94.8620μs 56.1168μs 17.8200 KOps/s 17.6922 KOps/s $\color{#35bf28}+0.72\%$
test_membership 4.6718μs 0.8413μs 1.1887 MOps/s 1.1906 MOps/s $\color{#d91a1a}-0.16\%$
test_membership_nested 30.3200μs 2.9981μs 333.5497 KOps/s 334.0803 KOps/s $\color{#d91a1a}-0.16\%$
test_membership_nested_leaf 33.4500μs 3.0165μs 331.5088 KOps/s 336.3406 KOps/s $\color{#d91a1a}-1.44\%$
test_membership_stacked_nested 34.8610μs 2.9185μs 342.6381 KOps/s 346.9457 KOps/s $\color{#d91a1a}-1.24\%$
test_membership_stacked_nested_leaf 31.4810μs 2.9185μs 342.6427 KOps/s 348.0283 KOps/s $\color{#d91a1a}-1.55\%$
test_membership_nested_last 37.6810μs 4.4787μs 223.2802 KOps/s 221.9058 KOps/s $\color{#35bf28}+0.62\%$
test_membership_nested_leaf_last 28.0910μs 4.5124μs 221.6140 KOps/s 222.0404 KOps/s $\color{#d91a1a}-0.19\%$
test_membership_stacked_nested_last 31.4410μs 4.3933μs 227.6185 KOps/s 229.1518 KOps/s $\color{#d91a1a}-0.67\%$
test_membership_stacked_nested_leaf_last 31.1610μs 4.3562μs 229.5601 KOps/s 228.0342 KOps/s $\color{#35bf28}+0.67\%$
test_nested_getleaf 53.5310μs 21.9412μs 45.5764 KOps/s 45.6542 KOps/s $\color{#d91a1a}-0.17\%$
test_nested_get 92.1520μs 20.8466μs 47.9694 KOps/s 48.4620 KOps/s $\color{#d91a1a}-1.02\%$
test_stacked_getleaf 73.3320μs 21.8433μs 45.7807 KOps/s 46.3459 KOps/s $\color{#d91a1a}-1.22\%$
test_stacked_get 43.1110μs 20.8631μs 47.9315 KOps/s 48.8857 KOps/s $\color{#d91a1a}-1.95\%$
test_nested_getitemleaf 0.2164ms 22.3177μs 44.8074 KOps/s 44.8322 KOps/s $\color{#d91a1a}-0.06\%$
test_nested_getitem 58.4810μs 21.0617μs 47.4797 KOps/s 46.9950 KOps/s $\color{#35bf28}+1.03\%$
test_stacked_getitemleaf 93.2620μs 22.3176μs 44.8076 KOps/s 45.0583 KOps/s $\color{#d91a1a}-0.56\%$
test_stacked_getitem 39.7000μs 21.4353μs 46.6521 KOps/s 47.3402 KOps/s $\color{#d91a1a}-1.45\%$
test_lock_nested 0.5746ms 0.4969ms 2.0125 KOps/s 1.9734 KOps/s $\color{#35bf28}+1.98\%$
test_lock_stack_nested 0.5563ms 0.5046ms 1.9818 KOps/s 1.9697 KOps/s $\color{#35bf28}+0.62\%$
test_unlock_nested 0.5040ms 0.3934ms 2.5416 KOps/s 2.5230 KOps/s $\color{#35bf28}+0.74\%$
test_unlock_stack_nested 0.5076ms 0.3981ms 2.5116 KOps/s 2.4947 KOps/s $\color{#35bf28}+0.68\%$
test_flatten_speed 0.1602ms 0.1218ms 8.2078 KOps/s 8.2169 KOps/s $\color{#d91a1a}-0.11\%$
test_unflatten_speed 0.6461ms 0.5709ms 1.7515 KOps/s 1.7246 KOps/s $\color{#35bf28}+1.56\%$
test_common_ops 0.1054s 1.7946ms 557.2158 Ops/s 243.4219 Ops/s $\textbf{\color{#35bf28}+128.91\%}$
test_creation 0.1225ms 3.1819μs 314.2801 KOps/s 318.0253 KOps/s $\color{#d91a1a}-1.18\%$
test_creation_empty 50.6510μs 6.9817μs 143.2311 KOps/s 141.7117 KOps/s $\color{#35bf28}+1.07\%$
test_creation_nested_1 41.8110μs 11.4843μs 87.0751 KOps/s 87.3079 KOps/s $\color{#d91a1a}-0.27\%$
test_creation_nested_2 46.1410μs 13.6035μs 73.5103 KOps/s 74.4729 KOps/s $\color{#d91a1a}-1.29\%$
test_creation_many_keys[10] 55.5410μs 21.1319μs 47.3219 KOps/s 48.1480 KOps/s $\color{#d91a1a}-1.72\%$
test_creation_many_keys[50] 0.1262ms 90.1789μs 11.0891 KOps/s 11.1132 KOps/s $\color{#d91a1a}-0.22\%$
test_creation_many_keys[100] 0.2255ms 0.1775ms 5.6348 KOps/s 5.6569 KOps/s $\color{#d91a1a}-0.39\%$
test_creation_nested_many_keys[10] 77.4710μs 45.3543μs 22.0486 KOps/s 22.3637 KOps/s $\color{#d91a1a}-1.41\%$
test_creation_nested_many_keys[50] 0.2397ms 0.1829ms 5.4660 KOps/s 5.6286 KOps/s $\color{#d91a1a}-2.89\%$
test_clone 40.2810μs 13.3380μs 74.9738 KOps/s 74.5591 KOps/s $\color{#35bf28}+0.56\%$
test_getitem[int] 1.4951ms 15.1241μs 66.1195 KOps/s 65.3567 KOps/s $\color{#35bf28}+1.17\%$
test_getitem[slice_int] 0.1447ms 24.7331μs 40.4317 KOps/s 40.2233 KOps/s $\color{#35bf28}+0.52\%$
test_getitem[range] 0.1803ms 66.0749μs 15.1344 KOps/s 14.5527 KOps/s $\color{#35bf28}+4.00\%$
test_getitem[tuple] 0.1489ms 23.9329μs 41.7834 KOps/s 41.1976 KOps/s $\color{#35bf28}+1.42\%$
test_getitem[list] 0.1860ms 58.8746μs 16.9852 KOps/s 16.7907 KOps/s $\color{#35bf28}+1.16\%$
test_setitem_dim[int] 46.2210μs 24.9443μs 40.0894 KOps/s 39.6509 KOps/s $\color{#35bf28}+1.11\%$
test_setitem_dim[slice_int] 68.3820μs 42.8333μs 23.3463 KOps/s 22.8223 KOps/s $\color{#35bf28}+2.30\%$
test_setitem_dim[range] 0.1315ms 97.7370μs 10.2315 KOps/s 10.0281 KOps/s $\color{#35bf28}+2.03\%$
test_setitem_dim[tuple] 74.1510μs 39.4532μs 25.3465 KOps/s 24.7219 KOps/s $\color{#35bf28}+2.53\%$
test_setitem 53.5210μs 17.6377μs 56.6967 KOps/s 55.2761 KOps/s $\color{#35bf28}+2.57\%$
test_set 0.1419ms 16.9424μs 59.0234 KOps/s 59.0323 KOps/s $\color{#d91a1a}-0.02\%$
test_set_shared 0.7046ms 0.2153ms 4.6439 KOps/s 4.6300 KOps/s $\color{#35bf28}+0.30\%$
test_update 0.2026ms 21.9943μs 45.4664 KOps/s 45.5867 KOps/s $\color{#d91a1a}-0.26\%$
test_update_nested 79.4020μs 33.4583μs 29.8879 KOps/s 29.1249 KOps/s $\color{#35bf28}+2.62\%$
test_update__nested 0.5419ms 32.5058μs 30.7637 KOps/s 30.0353 KOps/s $\color{#35bf28}+2.43\%$
test_set_nested 0.1015ms 18.7708μs 53.2743 KOps/s 49.8934 KOps/s $\textbf{\color{#35bf28}+6.78\%}$
test_set_nested_new 63.8920μs 23.8403μs 41.9458 KOps/s 40.3705 KOps/s $\color{#35bf28}+3.90\%$
test_select 0.4639ms 39.6195μs 25.2401 KOps/s 24.2635 KOps/s $\color{#35bf28}+4.03\%$
test_select_nested 0.5012ms 75.4969μs 13.2456 KOps/s 13.4350 KOps/s $\color{#d91a1a}-1.41\%$
test_exclude_nested 0.5172ms 92.1377μs 10.8533 KOps/s 10.8919 KOps/s $\color{#d91a1a}-0.35\%$
test_empty[True] 0.7422ms 0.4035ms 2.4782 KOps/s 2.4886 KOps/s $\color{#d91a1a}-0.42\%$
test_empty[False] 0.1054ms 1.3751μs 727.2193 KOps/s 719.9626 KOps/s $\color{#35bf28}+1.01\%$
test_to 0.5322ms 69.0872μs 14.4745 KOps/s 13.6314 KOps/s $\textbf{\color{#35bf28}+6.18\%}$
test_to_nonblocking 0.1183ms 69.0912μs 14.4736 KOps/s 13.9195 KOps/s $\color{#35bf28}+3.98\%$
test_unbind_speed 0.7822ms 0.3313ms 3.0188 KOps/s 3.0043 KOps/s $\color{#35bf28}+0.48\%$
test_unbind_speed_stack0 0.7724ms 0.3294ms 3.0354 KOps/s 3.0217 KOps/s $\color{#35bf28}+0.45\%$
test_unbind_speed_stack1 0.1133s 0.8491ms 1.1778 KOps/s 1.1634 KOps/s $\color{#35bf28}+1.24\%$
test_split 0.1144s 1.2795ms 781.5294 Ops/s 710.0299 Ops/s $\textbf{\color{#35bf28}+10.07\%}$
test_chunk 1.5098ms 1.0898ms 917.5683 Ops/s 921.0968 Ops/s $\color{#d91a1a}-0.38\%$
test_to_cpu_blocking 20.5188ms 20.1227ms 49.6951 Ops/s 49.6640 Ops/s $\color{#35bf28}+0.06\%$
test_to_cpu_global_sync 12.3779ms 11.7735ms 84.9364 Ops/s 85.1026 Ops/s $\color{#d91a1a}-0.20\%$
test_to_cpu_event_sync 13.6202ms 12.7684ms 78.3183 Ops/s 78.6760 Ops/s $\color{#d91a1a}-0.45\%$
test_to_cpu_default 13.2853ms 12.7646ms 78.3418 Ops/s 70.1062 Ops/s $\textbf{\color{#35bf28}+11.75\%}$
test_consolidate[False-None] 4.4045ms 4.2164ms 237.1706 Ops/s 237.3477 Ops/s $\color{#d91a1a}-0.07\%$
test_consolidate[default-None] 3.4714ms 2.1517ms 464.7389 Ops/s 456.4545 Ops/s $\color{#35bf28}+1.81\%$
test_consolidate[reduce-overhead-None] 3.4178ms 2.0759ms 481.7192 Ops/s 471.7726 Ops/s $\color{#35bf28}+2.11\%$
test_consolidate_njt[False-None] 8.7277ms 8.6783ms 115.2294 Ops/s 116.5835 Ops/s $\color{#d91a1a}-1.16\%$
test_to[False-False-None] 2.2443ms 2.1824ms 458.2156 Ops/s 456.6710 Ops/s $\color{#35bf28}+0.34\%$
test_to[True-False-None] 2.0538ms 1.9261ms 519.1880 Ops/s 527.5222 Ops/s $\color{#d91a1a}-1.58\%$
test_to[within-False-None] 6.6061ms 6.2751ms 159.3605 Ops/s 160.7642 Ops/s $\color{#d91a1a}-0.87\%$
test_to[True-default-None] 10.8377ms 10.6026ms 94.3165 Ops/s 92.4610 Ops/s $\color{#35bf28}+2.01\%$
test_to_njt[False-False-None] 8.7674ms 8.6020ms 116.2526 Ops/s 114.1792 Ops/s $\color{#35bf28}+1.82\%$
test_to_njt[True-False-None] 7.0590ms 6.8821ms 145.3046 Ops/s 140.5811 Ops/s $\color{#35bf28}+3.36\%$
test_to_njt[within-False-None] 15.8025ms 15.4344ms 64.7905 Ops/s 62.8698 Ops/s $\color{#35bf28}+3.06\%$
test_creation[device0] 0.3938ms 0.1168ms 8.5629 KOps/s 8.5133 KOps/s $\color{#35bf28}+0.58\%$
test_creation_from_tensor 0.4109ms 0.1141ms 8.7639 KOps/s 8.5168 KOps/s $\color{#35bf28}+2.90\%$
test_add_one[memmap_tensor0] 0.3199ms 6.7578μs 147.9762 KOps/s 146.9807 KOps/s $\color{#35bf28}+0.68\%$
test_contiguous[memmap_tensor0] 13.7900μs 0.7209μs 1.3872 MOps/s 2.0233 MOps/s $\textbf{\color{#d91a1a}-31.44\%}$
test_stack[memmap_tensor0] 35.5310μs 4.1843μs 238.9871 KOps/s 239.2193 KOps/s $\color{#d91a1a}-0.10\%$
test_memmaptd_index 1.1960ms 0.2785ms 3.5902 KOps/s 3.6175 KOps/s $\color{#d91a1a}-0.75\%$
test_memmaptd_index_astensor 0.5338ms 0.3817ms 2.6199 KOps/s 2.6004 KOps/s $\color{#35bf28}+0.75\%$
test_memmaptd_index_op 0.7662ms 0.6313ms 1.5839 KOps/s 1.4817 KOps/s $\textbf{\color{#35bf28}+6.90\%}$
test_serialize_model 0.1381s 0.1355s 7.3785 Ops/s 7.3490 Ops/s $\color{#35bf28}+0.40\%$
test_serialize_model_pickle 1.3645s 1.2121s 0.8250 Ops/s 0.8268 Ops/s $\color{#d91a1a}-0.22\%$
test_serialize_weights 0.1350s 0.1339s 7.4695 Ops/s 7.4584 Ops/s $\color{#35bf28}+0.15\%$
test_serialize_weights_returnearly 0.4416s 87.1235ms 11.4780 Ops/s 15.6367 Ops/s $\textbf{\color{#d91a1a}-26.60\%}$
test_serialize_weights_pickle 1.4230s 1.2223s 0.8182 Ops/s 0.8169 Ops/s $\color{#35bf28}+0.16\%$
test_reshape_pytree 0.1996ms 32.7405μs 30.5432 KOps/s 30.6020 KOps/s $\color{#d91a1a}-0.19\%$
test_reshape_td 0.1184ms 44.0328μs 22.7104 KOps/s 22.2326 KOps/s $\color{#35bf28}+2.15\%$
test_view_pytree 0.2139ms 32.1920μs 31.0636 KOps/s 31.0477 KOps/s $\color{#35bf28}+0.05\%$
test_view_td 0.1484ms 53.4453μs 18.7107 KOps/s 18.7767 KOps/s $\color{#d91a1a}-0.35\%$
test_unbind_pytree 0.2450ms 36.5838μs 27.3345 KOps/s 27.0242 KOps/s $\color{#35bf28}+1.15\%$
test_unbind_td 0.1674ms 49.4422μs 20.2256 KOps/s 19.7696 KOps/s $\color{#35bf28}+2.31\%$
test_split_pytree 0.1931ms 41.9236μs 23.8529 KOps/s 23.4240 KOps/s $\color{#35bf28}+1.83\%$
test_split_td 0.1856ms 64.0057μs 15.6236 KOps/s 15.3474 KOps/s $\color{#35bf28}+1.80\%$
test_add_pytree 0.2294ms 41.6103μs 24.0325 KOps/s 23.5095 KOps/s $\color{#35bf28}+2.22\%$
test_add_td 0.1296ms 54.5212μs 18.3415 KOps/s 17.7950 KOps/s $\color{#35bf28}+3.07\%$
test_compile_add_one_nested[tensordict-compile] 0.3412ms 0.2327ms 4.2979 KOps/s 4.3451 KOps/s $\color{#d91a1a}-1.09\%$
test_compile_add_one_nested[tensordict-eager] 0.3885ms 0.2041ms 4.8993 KOps/s 4.9959 KOps/s $\color{#d91a1a}-1.93\%$
test_compile_add_one_nested[pytree-compile] 0.2893ms 0.1353ms 7.3920 KOps/s 7.2179 KOps/s $\color{#35bf28}+2.41\%$
test_compile_add_one_nested[pytree-eager] 0.4572ms 0.1766ms 5.6620 KOps/s 5.6480 KOps/s $\color{#35bf28}+0.25\%$
test_compile_copy_nested[tensordict-compile] 0.1585ms 15.2054μs 65.7660 KOps/s 64.8147 KOps/s $\color{#35bf28}+1.47\%$
test_compile_copy_nested[tensordict-eager] 0.3329ms 53.3691μs 18.7374 KOps/s 18.4006 KOps/s $\color{#35bf28}+1.83\%$
test_compile_copy_nested[pytree-compile] 0.1263ms 15.9221μs 62.8058 KOps/s 64.2057 KOps/s $\color{#d91a1a}-2.18\%$
test_compile_copy_nested[pytree-eager] 0.3769ms 67.3879μs 14.8395 KOps/s 14.7934 KOps/s $\color{#35bf28}+0.31\%$
test_compile_add_one_flat[tensordict-compile] 0.5217ms 0.2504ms 3.9937 KOps/s 3.7512 KOps/s $\textbf{\color{#35bf28}+6.46\%}$
test_compile_add_one_flat[tensordict-eager] 0.4296ms 0.2822ms 3.5442 KOps/s 3.4668 KOps/s $\color{#35bf28}+2.23\%$
test_compile_add_one_flat[tensorclass-compile] 0.2910ms 0.1868ms 5.3523 KOps/s 5.0506 KOps/s $\textbf{\color{#35bf28}+5.97\%}$
test_compile_add_one_flat[tensorclass-eager] 0.2112ms 74.7161μs 13.3840 KOps/s 13.4034 KOps/s $\color{#d91a1a}-0.15\%$
test_compile_add_one_flat[pytree-compile] 0.4680ms 0.1931ms 5.1792 KOps/s 5.0294 KOps/s $\color{#35bf28}+2.98\%$
test_compile_add_one_flat[pytree-eager] 0.8074ms 0.5182ms 1.9298 KOps/s 1.9307 KOps/s $\color{#d91a1a}-0.05\%$
test_compile_add_self_flat[tensordict-eager] 0.4779ms 0.3311ms 3.0201 KOps/s 2.9977 KOps/s $\color{#35bf28}+0.75\%$
test_compile_add_self_flat[tensordict-compile] 0.4006ms 0.2519ms 3.9701 KOps/s 3.7585 KOps/s $\textbf{\color{#35bf28}+5.63\%}$
test_compile_add_self_flat[tensorclass-eager] 0.2260ms 91.3599μs 10.9457 KOps/s 11.0258 KOps/s $\color{#d91a1a}-0.73\%$
test_compile_add_self_flat[tensorclass-compile] 0.3383ms 0.1876ms 5.3312 KOps/s 5.0820 KOps/s $\color{#35bf28}+4.90\%$
test_compile_add_self_flat[pytree-eager] 0.6566ms 0.4330ms 2.3095 KOps/s 2.3294 KOps/s $\color{#d91a1a}-0.86\%$
test_compile_add_self_flat[pytree-compile] 0.3503ms 0.1934ms 5.1693 KOps/s 4.9609 KOps/s $\color{#35bf28}+4.20\%$
test_compile_copy_flat[tensordict-compile] 57.4910μs 19.3984μs 51.5505 KOps/s 50.1628 KOps/s $\color{#35bf28}+2.77\%$
test_compile_copy_flat[tensordict-eager] 72.4420μs 41.5489μs 24.0681 KOps/s 24.0749 KOps/s $\color{#d91a1a}-0.03\%$
test_compile_copy_flat[pytree-compile] 0.2153ms 20.9504μs 47.7318 KOps/s 45.2456 KOps/s $\textbf{\color{#35bf28}+5.49\%}$
test_compile_copy_flat[pytree-eager] 0.3259ms 52.2406μs 19.1422 KOps/s 19.2663 KOps/s $\color{#d91a1a}-0.64\%$
test_compile_assign_and_add[tensordict-compile] 2.1822ms 0.2389ms 4.1854 KOps/s 3.8221 KOps/s $\textbf{\color{#35bf28}+9.51\%}$
test_compile_assign_and_add[tensordict-eager] 3.9612ms 3.2278ms 309.8128 Ops/s 301.0135 Ops/s $\color{#35bf28}+2.92\%$
test_compile_assign_and_add[pytree-compile] 2.1067ms 0.1797ms 5.5642 KOps/s 5.4227 KOps/s $\color{#35bf28}+2.61\%$
test_compile_assign_and_add[pytree-eager] 3.5726ms 2.7577ms 362.6248 Ops/s 363.6280 Ops/s $\color{#d91a1a}-0.28\%$
test_compile_indexing[tensor-tensordict-compile] 0.2860ms 0.1344ms 7.4402 KOps/s 7.2838 KOps/s $\color{#35bf28}+2.15\%$
test_compile_indexing[tensor-tensordict-eager] 0.3097ms 74.5852μs 13.4075 KOps/s 13.4036 KOps/s $\color{#35bf28}+0.03\%$
test_compile_indexing[tensor-tensorclass-compile] 0.2397ms 0.1195ms 8.3661 KOps/s 8.3734 KOps/s $\color{#d91a1a}-0.09\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2508ms 43.8664μs 22.7965 KOps/s 22.3961 KOps/s $\color{#35bf28}+1.79\%$
test_compile_indexing[tensor-pytree-compile] 0.1618ms 0.1212ms 8.2530 KOps/s 8.2811 KOps/s $\color{#d91a1a}-0.34\%$
test_compile_indexing[tensor-pytree-eager] 0.2491ms 43.8897μs 22.7844 KOps/s 22.4501 KOps/s $\color{#35bf28}+1.49\%$
test_compile_indexing[slice-tensordict-compile] 0.1974ms 82.7076μs 12.0908 KOps/s 11.9441 KOps/s $\color{#35bf28}+1.23\%$
test_compile_indexing[slice-tensordict-eager] 0.2103ms 27.2988μs 36.6317 KOps/s 36.6727 KOps/s $\color{#d91a1a}-0.11\%$
test_compile_indexing[slice-tensorclass-compile] 0.1920ms 69.4807μs 14.3925 KOps/s 14.4101 KOps/s $\color{#d91a1a}-0.12\%$
test_compile_indexing[slice-tensorclass-eager] 0.2392ms 21.9619μs 45.5333 KOps/s 45.4641 KOps/s $\color{#35bf28}+0.15\%$
test_compile_indexing[slice-pytree-compile] 0.1153ms 70.2827μs 14.2283 KOps/s 14.1133 KOps/s $\color{#35bf28}+0.81\%$
test_compile_indexing[slice-pytree-eager] 0.2595ms 21.8782μs 45.7075 KOps/s 45.4891 KOps/s $\color{#35bf28}+0.48\%$
test_compile_indexing[int-tensordict-compile] 0.1297ms 83.1069μs 12.0327 KOps/s 11.9011 KOps/s $\color{#35bf28}+1.11\%$
test_compile_indexing[int-tensordict-eager] 0.2064ms 26.8945μs 37.1824 KOps/s 36.4616 KOps/s $\color{#35bf28}+1.98\%$
test_compile_indexing[int-tensorclass-compile] 0.1953ms 69.7060μs 14.3460 KOps/s 14.2226 KOps/s $\color{#35bf28}+0.87\%$
test_compile_indexing[int-tensorclass-eager] 0.2252ms 21.7259μs 46.0280 KOps/s 45.5650 KOps/s $\color{#35bf28}+1.02\%$
test_compile_indexing[int-pytree-compile] 0.1067ms 69.9081μs 14.3045 KOps/s 14.2115 KOps/s $\color{#35bf28}+0.65\%$
test_compile_indexing[int-pytree-eager] 0.2583ms 21.6900μs 46.1042 KOps/s 45.4520 KOps/s $\color{#35bf28}+1.43\%$
test_compile_replace[single-eager] 0.2331ms 49.4038μs 20.2414 KOps/s 21.1726 KOps/s $\color{#d91a1a}-4.40\%$
test_compile_replace[single-compile] 0.2319ms 0.1331ms 7.5157 KOps/s 7.4528 KOps/s $\color{#35bf28}+0.84\%$
test_compile_replace[multi-eager] 0.6535ms 0.5798ms 1.7247 KOps/s 1.7937 KOps/s $\color{#d91a1a}-3.85\%$
test_compile_replace[multi-compile] 0.3084ms 0.1418ms 7.0519 KOps/s 7.0474 KOps/s $\color{#35bf28}+0.06\%$
test_compile_tc_getattr_20[eager] 0.2132ms 0.1670ms 5.9879 KOps/s 6.2543 KOps/s $\color{#d91a1a}-4.26\%$
test_compile_tc_getattr_20[compile] 0.3764ms 0.1530ms 6.5373 KOps/s 6.6437 KOps/s $\color{#d91a1a}-1.60\%$
test_compile_clone_shallow[20-eager] 57.2210μs 19.0678μs 52.4445 KOps/s 52.0040 KOps/s $\color{#35bf28}+0.85\%$
test_compile_clone_shallow[20-compile] 81.0320μs 16.7424μs 59.7287 KOps/s 57.7805 KOps/s $\color{#35bf28}+3.37\%$
test_compile_clone_shallow[40-eager] 69.2910μs 33.6771μs 29.6937 KOps/s 29.2008 KOps/s $\color{#35bf28}+1.69\%$
test_compile_clone_shallow[40-compile] 53.8410μs 18.9993μs 52.6336 KOps/s 52.0234 KOps/s $\color{#35bf28}+1.17\%$
test_compile_clone_shallow[80-eager] 93.1610μs 62.8372μs 15.9141 KOps/s 16.0825 KOps/s $\color{#d91a1a}-1.05\%$
test_compile_clone_shallow[80-compile] 79.0920μs 21.9472μs 45.5638 KOps/s 46.0617 KOps/s $\color{#d91a1a}-1.08\%$
test_compile_update_inplace[eager] 0.1800ms 60.2473μs 16.5983 KOps/s 16.4240 KOps/s $\color{#35bf28}+1.06\%$
test_compile_update_inplace[compile] 0.4019ms 0.1462ms 6.8386 KOps/s 6.7037 KOps/s $\color{#35bf28}+2.01\%$
test_mod_add[eager] 0.4532ms 0.2267ms 4.4119 KOps/s 4.4123 KOps/s $\color{#d91a1a}-0.01\%$
test_mod_add[compile] 0.2904ms 0.1352ms 7.3985 KOps/s 7.3680 KOps/s $\color{#35bf28}+0.41\%$
test_mod_add[compile-overhead] 0.3240ms 0.1793ms 5.5783 KOps/s 5.5127 KOps/s $\color{#35bf28}+1.19\%$
test_mod_wrap[eager] 0.5856ms 0.4675ms 2.1390 KOps/s 2.1294 KOps/s $\color{#35bf28}+0.45\%$
test_mod_wrap[compile] 0.5338ms 0.3928ms 2.5458 KOps/s 2.4852 KOps/s $\color{#35bf28}+2.44\%$
test_mod_wrap[compile-overhead] 8.4731ms 4.6481ms 215.1438 Ops/s 216.7977 Ops/s $\color{#d91a1a}-0.76\%$
test_mod_wrap_and_backward[eager] 2.0100ms 1.7962ms 556.7374 Ops/s 586.5838 Ops/s $\textbf{\color{#d91a1a}-5.09\%}$
test_mod_wrap_and_backward[compile] 1.8151ms 1.6407ms 609.4857 Ops/s 643.5788 Ops/s $\textbf{\color{#d91a1a}-5.30\%}$
test_mod_wrap_and_backward[compile-overhead] 1.3385ms 0.9468ms 1.0562 KOps/s 1.0394 KOps/s $\color{#35bf28}+1.62\%$
test_seq_add[eager] 0.6637ms 0.5108ms 1.9577 KOps/s 1.9640 KOps/s $\color{#d91a1a}-0.32\%$
test_seq_add[compile] 0.2853ms 0.1436ms 6.9635 KOps/s 6.8345 KOps/s $\color{#35bf28}+1.89\%$
test_seq_add[compile-overhead] 0.3148ms 0.1841ms 5.4330 KOps/s 5.2752 KOps/s $\color{#35bf28}+2.99\%$
test_seq_wrap[eager] 1.5603ms 1.4134ms 707.5110 Ops/s 692.6808 Ops/s $\color{#35bf28}+2.14\%$
test_seq_wrap[compile] 0.5837ms 0.4133ms 2.4193 KOps/s 2.2844 KOps/s $\textbf{\color{#35bf28}+5.90\%}$
test_seq_wrap[compile-overhead] 0.4028ms 0.3045ms 3.2838 KOps/s 3.1841 KOps/s $\color{#35bf28}+3.13\%$
test_func_call_runtime[False-eager] 0.9418ms 0.8219ms 1.2167 KOps/s 1.2583 KOps/s $\color{#d91a1a}-3.30\%$
test_func_call_runtime[False-compile] 1.1966ms 1.0004ms 999.6133 Ops/s 972.2340 Ops/s $\color{#35bf28}+2.82\%$
test_func_call_runtime[False-compile-overhead] 0.6498ms 0.5150ms 1.9416 KOps/s 1.9311 KOps/s $\color{#35bf28}+0.55\%$
test_func_call_runtime[True-eager] 1.1849ms 1.0488ms 953.4942 Ops/s 948.8827 Ops/s $\color{#35bf28}+0.49\%$
test_func_call_runtime[True-compile] 1.1313ms 0.9985ms 1.0015 KOps/s 964.4383 Ops/s $\color{#35bf28}+3.85\%$
test_func_call_runtime[True-compile-overhead] 0.5729ms 0.5248ms 1.9054 KOps/s 1.8642 KOps/s $\color{#35bf28}+2.21\%$
test_func_call_cm_runtime[False-eager] 0.9184ms 0.8276ms 1.2083 KOps/s 1.2254 KOps/s $\color{#d91a1a}-1.39\%$
test_func_call_cm_runtime[False-compile] 1.3618ms 0.9750ms 1.0257 KOps/s 1.0249 KOps/s $\color{#35bf28}+0.07\%$
test_func_call_cm_runtime[False-compile-overhead] 0.7765ms 0.4985ms 2.0061 KOps/s 2.0040 KOps/s $\color{#35bf28}+0.10\%$
test_func_call_cm_runtime[True-eager] 1.3519ms 1.1961ms 836.0206 Ops/s 831.8654 Ops/s $\color{#35bf28}+0.50\%$
test_func_call_cm_runtime[True-compile] 1.1669ms 1.0325ms 968.5058 Ops/s 968.3351 Ops/s $\color{#35bf28}+0.02\%$
test_func_call_cm_runtime[True-compile-overhead] 0.7170ms 0.5505ms 1.8166 KOps/s 1.7956 KOps/s $\color{#35bf28}+1.17\%$
test_vmap_func_call_cm_runtime[eager] 2.8665ms 2.4050ms 415.8062 Ops/s 414.1467 Ops/s $\color{#35bf28}+0.40\%$
test_vmap_func_call_cm_runtime[compile] 3.0453ms 2.5926ms 385.7071 Ops/s 383.7482 Ops/s $\color{#35bf28}+0.51\%$
test_vmap_func_call_cm_runtime[compile-overhead] 3.3331ms 2.5975ms 384.9866 Ops/s 383.7201 Ops/s $\color{#35bf28}+0.33\%$
test_distributed 2.8771ms 0.1575ms 6.3476 KOps/s 6.5202 KOps/s $\color{#d91a1a}-2.65\%$
test_tdmodule 0.4472ms 0.1921ms 5.2051 KOps/s 5.0594 KOps/s $\color{#35bf28}+2.88\%$
test_tdmodule_dispatch 0.4251ms 0.3724ms 2.6852 KOps/s 2.6023 KOps/s $\color{#35bf28}+3.18\%$
test_tdseq 47.2510μs 26.5548μs 37.6580 KOps/s 35.8061 KOps/s $\textbf{\color{#35bf28}+5.17\%}$
test_tdseq_dispatch 0.2475ms 0.2160ms 4.6293 KOps/s 4.6045 KOps/s $\color{#35bf28}+0.54\%$
test_instantiation_functorch 2.1651ms 2.0640ms 484.4985 Ops/s 485.2025 Ops/s $\color{#d91a1a}-0.15\%$
test_exec_functorch 0.2427ms 0.1805ms 5.5414 KOps/s 5.4898 KOps/s $\color{#35bf28}+0.94\%$
test_exec_functional_call 0.2046ms 0.1575ms 6.3485 KOps/s 6.2762 KOps/s $\color{#35bf28}+1.15\%$
test_exec_td_decorator 0.4906ms 0.2497ms 4.0047 KOps/s 4.0063 KOps/s $\color{#d91a1a}-0.04\%$
test_vmap_mlp_speed_decorator[True-True] 1.2270ms 1.0318ms 969.1361 Ops/s 963.7680 Ops/s $\color{#35bf28}+0.56\%$
test_vmap_mlp_speed_decorator[True-False] 1.2018ms 1.0303ms 970.6031 Ops/s 964.7825 Ops/s $\color{#35bf28}+0.60\%$
test_vmap_mlp_speed_decorator[False-True] 0.9357ms 0.7306ms 1.3687 KOps/s 1.3504 KOps/s $\color{#35bf28}+1.35\%$
test_vmap_mlp_speed_decorator[False-False] 0.9411ms 0.7319ms 1.3662 KOps/s 1.3621 KOps/s $\color{#35bf28}+0.30\%$
test_vmap_transformer_speed_decorator[True-True] 21.9401ms 21.0378ms 47.5336 Ops/s 47.3772 Ops/s $\color{#35bf28}+0.33\%$
test_vmap_transformer_speed_decorator[True-False] 21.7630ms 21.0300ms 47.5511 Ops/s 47.3896 Ops/s $\color{#35bf28}+0.34\%$
test_vmap_transformer_speed_decorator[False-True] 20.8704ms 20.6264ms 48.4816 Ops/s 48.0599 Ops/s $\color{#35bf28}+0.88\%$
test_vmap_transformer_speed_decorator[False-False] 20.8586ms 20.5952ms 48.5551 Ops/s 48.2155 Ops/s $\color{#35bf28}+0.70\%$
test_to_module_speed[True] 1.7790ms 1.6296ms 613.6454 Ops/s 606.2163 Ops/s $\color{#35bf28}+1.23\%$
test_to_module_speed[False] 1.6876ms 1.6068ms 622.3689 Ops/s 619.9087 Ops/s $\color{#35bf28}+0.40\%$
test_to_module_plain_tensor_speed[None] 1.8398ms 1.7333ms 576.9423 Ops/s 574.1966 Ops/s $\color{#35bf28}+0.48\%$
test_to_module_plain_tensor_speed[False] 3.0078ms 1.5658ms 638.6335 Ops/s 639.7845 Ops/s $\color{#d91a1a}-0.18\%$
test_to_module_plain_tensor_speed[True] 4.0919ms 3.4497ms 289.8805 Ops/s 291.3102 Ops/s $\color{#d91a1a}-0.49\%$
test_tc_init 83.0320μs 43.6951μs 22.8859 KOps/s 23.4918 KOps/s $\color{#d91a1a}-2.58\%$
test_tc_init_tensor_only 28.0800μs 9.1140μs 109.7216 KOps/s 110.4727 KOps/s $\color{#d91a1a}-0.68\%$
test_tc_init_nested 0.1344ms 86.7385μs 11.5289 KOps/s 11.6331 KOps/s $\color{#d91a1a}-0.90\%$
test_tc_init_many_fields 45.2510μs 15.1402μs 66.0493 KOps/s 67.7791 KOps/s $\color{#d91a1a}-2.55\%$
test_tc_first_layer_tensor 66.1610μs 1.8144μs 551.1348 KOps/s 553.4671 KOps/s $\color{#d91a1a}-0.42\%$
test_tc_first_layer_tensor_only 1.8305μs 0.3941μs 2.5371 MOps/s 2.4978 MOps/s $\color{#35bf28}+1.57\%$
test_tc_first_layer_tensor_set 28.7200μs 3.9206μs 255.0632 KOps/s 256.5138 KOps/s $\color{#d91a1a}-0.57\%$
test_tc_first_layer_tensor_only_set 21.0300μs 3.2871μs 304.2204 KOps/s 302.7705 KOps/s $\color{#35bf28}+0.48\%$
test_tc_first_layer_nontensor 33.0810μs 6.1934μs 161.4610 KOps/s 157.4588 KOps/s $\color{#35bf28}+2.54\%$
test_tc_second_layer_tensor 21.4300μs 4.4481μs 224.8154 KOps/s 227.7343 KOps/s $\color{#d91a1a}-1.28\%$
test_tc_second_layer_nontensor 48.6010μs 8.5771μs 116.5890 KOps/s 113.2604 KOps/s $\color{#35bf28}+2.94\%$
test_unbind 0.2652s 13.5315ms 73.9019 Ops/s 65.4197 Ops/s $\textbf{\color{#35bf28}+12.97\%}$
test_full_like 4.8258ms 4.3748ms 228.5819 Ops/s 227.1062 Ops/s $\color{#35bf28}+0.65\%$
test_zeros_like 4.7485ms 4.3629ms 229.2047 Ops/s 227.9018 Ops/s $\color{#35bf28}+0.57\%$
test_ones_like 4.8454ms 4.3738ms 228.6319 Ops/s 227.6750 Ops/s $\color{#35bf28}+0.42\%$
test_clone 6.9503ms 6.5543ms 152.5711 Ops/s 147.3941 Ops/s $\color{#35bf28}+3.51\%$
test_squeeze 63.2220μs 14.4590μs 69.1611 KOps/s 70.8865 KOps/s $\color{#d91a1a}-2.43\%$
test_unsqueeze 0.2778ms 0.1115ms 8.9687 KOps/s 8.7222 KOps/s $\color{#35bf28}+2.83\%$
test_split 0.2707ms 0.1817ms 5.5030 KOps/s 5.3270 KOps/s $\color{#35bf28}+3.30\%$
test_permute 0.2689ms 0.2078ms 4.8134 KOps/s 4.6072 KOps/s $\color{#35bf28}+4.48\%$
test_stack 27.5984ms 26.9321ms 37.1304 Ops/s 53.4872 Ops/s $\textbf{\color{#d91a1a}-30.58\%}$
test_cat 27.6000ms 26.9803ms 37.0641 Ops/s 53.7759 Ops/s $\textbf{\color{#d91a1a}-31.08\%}$
test_sequential_tensordict 0.8364ms 0.5675ms 1.7622 KOps/s 1.6849 KOps/s $\color{#35bf28}+4.59\%$
test_sequential_graph_module 0.5398ms 0.1166ms 8.5767 KOps/s 8.4222 KOps/s $\color{#35bf28}+1.83\%$
test_nested_tensordict 0.9199ms 0.8136ms 1.2291 KOps/s 1.2185 KOps/s $\color{#35bf28}+0.87\%$
test_nested_graph_module 0.5437ms 0.1283ms 7.7920 KOps/s 7.6798 KOps/s $\color{#35bf28}+1.46\%$

@vmoens
vmoens force-pushed the fix-memmap-tensorclass branch from 49d905c to 3d55508 Compare June 10, 2026 17:39
…sorClass subclasses

TensorClass subclasses (e.g. class Foo(TensorClass["nocast"])) lost their
class identity through a memmap_ / load_memmap roundtrip, coming back as a
plain TensorDict holding the tensorclass fields. Decorator-based
tensorclasses (@Tensorclass) were unaffected at the top level but their
subclasses hit the same issue on load.

Two wiring bugs:

- "_memmap_" was listed in _FALLBACK_METHOD_FROM_TD, so the fallback loop in
  _patch_tc overwrote the dedicated tensorclass _memmap_ (which records the
  class in meta.json and nests the data under _tensordict/) with a
  TD-delegating wrapper. The node was saved as a plain TensorDict, leaving
  load_memmap nothing to dispatch on.

- The classmethod-wrapping loop in _tensorclass only checked cls.__dict__
  when deciding whether a TensorDict classmethod was "missing", shadowing
  the genuine _load_memmap classmethod inherited from the TensorClass base
  with TensorDict._load_memmap, which returns an empty instance on the
  tensorclass directory layout. The loop now skips attributes inherited as
  genuine classmethod/staticmethod implementations and keeps re-wrapping
  everything else (generated wrappers and parent-bound methods such as
  from_list or stack on the TensorClass base) per subclass, preserving the
  previous behavior of those entry points.

Directories saved with the old flat layout still load, degrading to a plain
TensorDict as before.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@vmoens
vmoens force-pushed the fix-memmap-tensorclass branch from 3d55508 to 3bc9abe Compare June 10, 2026 20:19
@vmoens
vmoens merged commit 3180e59 into main Jun 11, 2026
70 of 73 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working CI CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. tensorclass Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant