[Feature] as_decorator and __enter__ __exit__ methods

[vmoens]

No description provided.

[github-actions]

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

Total Benchmarks: 92. Improved: $\large\color{#35bf28}4$. Worsened: $\large\color{#d91a1a}12$.

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_items	25.2010μs	3.4461μs	290.1856 KOps/s	293.1561 KOps/s	$\color{#d91a1a}-1.01\%$
test_items_nested	1.7314ms	0.5659ms	1.7670 KOps/s	1.6918 KOps/s	$\color{#35bf28}+4.44\%$
test_items_nested_locked	0.2861ms	0.1949ms	5.1318 KOps/s	4.9192 KOps/s	$\color{#35bf28}+4.32\%$
test_items_nested_leaf	0.5431ms	0.3268ms	3.0603 KOps/s	2.7726 KOps/s	$\textbf{\color{#35bf28}+10.38\%}$
test_items_stack_nested	30.2932ms	30.0506ms	33.2773 Ops/s	33.2604 Ops/s	$\color{#35bf28}+0.05\%$
test_items_stack_nested_leaf	16.5088ms	16.3847ms	61.0324 Ops/s	60.8984 Ops/s	$\color{#35bf28}+0.22\%$
test_items_stack_nested_locked	1.7785ms	1.6603ms	602.3166 Ops/s	572.4739 Ops/s	$\textbf{\color{#35bf28}+5.21\%}$
test_keys	26.2010μs	5.0729μs	197.1268 KOps/s	196.6979 KOps/s	$\color{#35bf28}+0.22\%$
test_keys_nested	0.2373ms	0.1800ms	5.5554 KOps/s	5.5438 KOps/s	$\color{#35bf28}+0.21\%$
test_keys_nested_locked	0.2016ms	0.1796ms	5.5689 KOps/s	5.5409 KOps/s	$\color{#35bf28}+0.50\%$
test_keys_nested_leaf	0.1980ms	0.1775ms	5.6335 KOps/s	5.6221 KOps/s	$\color{#35bf28}+0.20\%$
test_keys_stack_nested	2.1732ms	2.0608ms	485.2384 Ops/s	480.8078 Ops/s	$\color{#35bf28}+0.92\%$
test_keys_stack_nested_leaf	2.1137ms	2.0609ms	485.2199 Ops/s	479.4226 Ops/s	$\color{#35bf28}+1.21\%$
test_keys_stack_nested_locked	0.7535ms	0.6499ms	1.5388 KOps/s	1.4778 KOps/s	$\color{#35bf28}+4.13\%$
test_values	22.4000μs	1.5558μs	642.7609 KOps/s	642.7739 KOps/s	$-0.00\%$
test_values_nested	0.6423ms	0.5630ms	1.7760 KOps/s	1.6954 KOps/s	$\color{#35bf28}+4.76\%$
test_values_nested_locked	0.2231ms	0.1929ms	5.1842 KOps/s	5.0027 KOps/s	$\color{#35bf28}+3.63\%$
test_values_nested_leaf	0.3560ms	0.3244ms	3.0827 KOps/s	2.9508 KOps/s	$\color{#35bf28}+4.47\%$
test_values_stack_nested	33.5518ms	30.0623ms	33.2643 Ops/s	33.2420 Ops/s	$\color{#35bf28}+0.07\%$
test_values_stack_nested_leaf	17.1540ms	16.3198ms	61.2754 Ops/s	60.7316 Ops/s	$\color{#35bf28}+0.90\%$
test_values_stack_nested_locked	1.7085ms	1.6587ms	602.8836 Ops/s	579.1291 Ops/s	$\color{#35bf28}+4.10\%$
test_membership	12.5000μs	1.8597μs	537.7195 KOps/s	537.5467 KOps/s	$\color{#35bf28}+0.03\%$
test_membership_nested	20.1000μs	3.3085μs	302.2513 KOps/s	295.7972 KOps/s	$\color{#35bf28}+2.18\%$
test_membership_nested_leaf	25.6010μs	3.3387μs	299.5203 KOps/s	297.6402 KOps/s	$\color{#35bf28}+0.63\%$
test_membership_stacked_nested	14.4000μs	1.9528μs	512.0910 KOps/s	527.1501 KOps/s	$\color{#d91a1a}-2.86\%$
test_membership_stacked_nested_leaf	27.7010μs	1.9016μs	525.8797 KOps/s	522.0970 KOps/s	$\color{#35bf28}+0.72\%$
test_stacked_getleaf	1.4290ms	1.3938ms	717.4421 Ops/s	718.3294 Ops/s	$\color{#d91a1a}-0.12\%$
test_stacked_get	1.4609ms	1.3527ms	739.2718 Ops/s	745.3822 Ops/s	$\color{#d91a1a}-0.82\%$
test_lock_nested	1.1755ms	0.9869ms	1.0133 KOps/s	1.1012 KOps/s	$\textbf{\color{#d91a1a}-7.99\%}$
test_lock_stack_nested	87.2404ms	13.6137ms	73.4555 Ops/s	75.2836 Ops/s	$\color{#d91a1a}-2.43\%$
test_unlock_nested	1.1736ms	0.9810ms	1.0193 KOps/s	1.0814 KOps/s	$\textbf{\color{#d91a1a}-5.74\%}$
test_unlock_stack_nested	86.1782ms	14.0105ms	71.3749 Ops/s	96.5962 Ops/s	$\textbf{\color{#d91a1a}-26.11\%}$
test_flatten_speed	0.9986ms	0.9504ms	1.0522 KOps/s	1.0678 KOps/s	$\color{#d91a1a}-1.46\%$
test_unflatten_speed	1.8281ms	1.7219ms	580.7464 Ops/s	592.9362 Ops/s	$\color{#d91a1a}-2.06\%$
test_common_ops	1.6046ms	1.1061ms	904.0437 Ops/s	942.0684 Ops/s	$\color{#d91a1a}-4.04\%$
test_creation	29.0000μs	6.4892μs	154.1016 KOps/s	173.4451 KOps/s	$\textbf{\color{#d91a1a}-11.15\%}$
test_creation_empty	31.5000μs	13.3711μs	74.7881 KOps/s	77.8769 KOps/s	$\color{#d91a1a}-3.97\%$
test_creation_nested_1	43.7010μs	26.2905μs	38.0365 KOps/s	39.9866 KOps/s	$\color{#d91a1a}-4.88\%$
test_creation_nested_2	50.6010μs	26.1002μs	38.3139 KOps/s	40.0336 KOps/s	$\color{#d91a1a}-4.30\%$
test_clone	0.2965ms	25.7045μs	38.9036 KOps/s	41.7332 KOps/s	$\textbf{\color{#d91a1a}-6.78\%}$
test_getitem[int]	99.2010μs	31.1812μs	32.0706 KOps/s	33.6150 KOps/s	$\color{#d91a1a}-4.59\%$
test_getitem[slice_int]	92.0020μs	67.5718μs	14.7991 KOps/s	15.3658 KOps/s	$\color{#d91a1a}-3.69\%$
test_getitem[range]	97.0010μs	69.1109μs	14.4695 KOps/s	15.2767 KOps/s	$\textbf{\color{#d91a1a}-5.28\%}$
test_getitem[tuple]	81.9010μs	62.9620μs	15.8826 KOps/s	16.6176 KOps/s	$\color{#d91a1a}-4.42\%$
test_getitem[list]	82.7010μs	60.2051μs	16.6099 KOps/s	17.3696 KOps/s	$\color{#d91a1a}-4.37\%$
test_setitem_dim[int]	59.6010μs	44.1239μs	22.6635 KOps/s	22.7735 KOps/s	$\color{#d91a1a}-0.48\%$
test_setitem_dim[slice_int]	0.1155ms	83.3333μs	12.0000 KOps/s	12.0583 KOps/s	$\color{#d91a1a}-0.48\%$
test_setitem_dim[range]	0.1010ms	77.7827μs	12.8563 KOps/s	12.9822 KOps/s	$\color{#d91a1a}-0.97\%$
test_setitem_dim[tuple]	95.4010μs	75.8602μs	13.1821 KOps/s	13.3684 KOps/s	$\color{#d91a1a}-1.39\%$
test_setitem	0.2553ms	32.6378μs	30.6393 KOps/s	32.6494 KOps/s	$\textbf{\color{#d91a1a}-6.16\%}$
test_set	0.2143ms	32.1834μs	31.0719 KOps/s	33.4629 KOps/s	$\textbf{\color{#d91a1a}-7.15\%}$
test_set_shared	0.2971ms	0.1739ms	5.7512 KOps/s	6.0372 KOps/s	$\color{#d91a1a}-4.74\%$
test_update	0.2173ms	34.1368μs	29.2939 KOps/s	29.5932 KOps/s	$\color{#d91a1a}-1.01\%$
test_update_nested	0.3339ms	51.4454μs	19.4381 KOps/s	20.2159 KOps/s	$\color{#d91a1a}-3.85\%$
test_set_nested	0.2356ms	43.3711μs	23.0568 KOps/s	24.2379 KOps/s	$\color{#d91a1a}-4.87\%$
test_set_nested_new	86.9010μs	62.4748μs	16.0064 KOps/s	16.8590 KOps/s	$\textbf{\color{#d91a1a}-5.06\%}$
test_select	0.1685ms	0.1060ms	9.4334 KOps/s	10.0422 KOps/s	$\textbf{\color{#d91a1a}-6.06\%}$
test_creation[device0]	0.5577ms	0.4439ms	2.2530 KOps/s	2.2576 KOps/s	$\color{#d91a1a}-0.21\%$
test_creation_from_tensor	0.5964ms	0.4954ms	2.0187 KOps/s	1.9748 KOps/s	$\color{#35bf28}+2.22\%$
test_add_one[memmap_tensor0]	0.2358ms	33.0835μs	30.2266 KOps/s	30.3133 KOps/s	$\color{#d91a1a}-0.29\%$
test_contiguous[memmap_tensor0]	31.4000μs	8.8472μs	113.0298 KOps/s	115.4998 KOps/s	$\color{#d91a1a}-2.14\%$
test_stack[memmap_tensor0]	0.1739ms	26.7663μs	37.3604 KOps/s	37.8937 KOps/s	$\color{#d91a1a}-1.41\%$
test_memmaptd_index	0.3276ms	0.2834ms	3.5287 KOps/s	3.5234 KOps/s	$\color{#35bf28}+0.15\%$
test_memmaptd_index_astensor	1.1080ms	1.0692ms	935.3119 Ops/s	921.0600 Ops/s	$\color{#35bf28}+1.55\%$
test_memmaptd_index_op	2.3820ms	2.2634ms	441.8119 Ops/s	445.3997 Ops/s	$\color{#d91a1a}-0.81\%$
test_reshape_pytree	95.1020μs	37.7345μs	26.5009 KOps/s	26.2792 KOps/s	$\color{#35bf28}+0.84\%$
test_reshape_td	74.3010μs	46.3088μs	21.5942 KOps/s	22.6767 KOps/s	$\color{#d91a1a}-4.77\%$
test_view_pytree	94.0010μs	35.3428μs	28.2943 KOps/s	28.2520 KOps/s	$\color{#35bf28}+0.15\%$
test_view_td	62.5010μs	9.7172μs	102.9101 KOps/s	99.9864 KOps/s	$\color{#35bf28}+2.92\%$
test_unbind_pytree	83.7010μs	38.9299μs	25.6872 KOps/s	25.1282 KOps/s	$\color{#35bf28}+2.22\%$
test_unbind_td	0.1673ms	0.1381ms	7.2423 KOps/s	7.5269 KOps/s	$\color{#d91a1a}-3.78\%$
test_split_pytree	0.1008ms	45.0805μs	22.1826 KOps/s	22.0680 KOps/s	$\color{#35bf28}+0.52\%$
test_split_td	0.9116ms	0.1164ms	8.5921 KOps/s	8.8909 KOps/s	$\color{#d91a1a}-3.36\%$
test_add_pytree	94.6010μs	48.5826μs	20.5835 KOps/s	20.8607 KOps/s	$\color{#d91a1a}-1.33\%$
test_add_td	81.4020μs	59.7993μs	16.7226 KOps/s	17.3790 KOps/s	$\color{#d91a1a}-3.78\%$
test_distributed	23.2000μs	8.7875μs	113.7974 KOps/s	109.0366 KOps/s	$\color{#35bf28}+4.37\%$
test_tdmodule	4.3501ms	30.2009μs	33.1116 KOps/s	37.8656 KOps/s	$\textbf{\color{#d91a1a}-12.55\%}$
test_tdmodule_dispatch	0.2741ms	56.0849μs	17.8301 KOps/s	17.6750 KOps/s	$\color{#35bf28}+0.88\%$
test_tdseq	0.1728ms	35.7823μs	27.9467 KOps/s	30.8458 KOps/s	$\textbf{\color{#d91a1a}-9.40\%}$
test_tdseq_dispatch	0.2079ms	67.2267μs	14.8750 KOps/s	14.8877 KOps/s	$\color{#d91a1a}-0.09\%$
test_instantiation_functorch	1.7512ms	1.6114ms	620.5944 Ops/s	615.1729 Ops/s	$\color{#35bf28}+0.88\%$
test_instantiation_td	1.9388ms	1.3460ms	742.9460 Ops/s	696.9036 Ops/s	$\textbf{\color{#35bf28}+6.61\%}$
test_exec_functorch	0.2215ms	0.1886ms	5.3028 KOps/s	5.2743 KOps/s	$\color{#35bf28}+0.54\%$
test_exec_td	0.2498ms	0.1755ms	5.6976 KOps/s	5.7629 KOps/s	$\color{#d91a1a}-1.13\%$
test_vmap_mlp_speed[True-True]	1.2398ms	0.8481ms	1.1791 KOps/s	1.1860 KOps/s	$\color{#d91a1a}-0.59\%$
test_vmap_mlp_speed[True-False]	1.8984ms	0.5729ms	1.7455 KOps/s	1.6169 KOps/s	$\textbf{\color{#35bf28}+7.95\%}$
test_vmap_mlp_speed[False-True]	1.2044ms	0.7030ms	1.4224 KOps/s	1.4429 KOps/s	$\color{#d91a1a}-1.42\%$
test_vmap_mlp_speed[False-False]	0.8914ms	0.4241ms	2.3582 KOps/s	2.3384 KOps/s	$\color{#35bf28}+0.85\%$
test_vmap_transformer_speed[True-True]	11.6660ms	10.5900ms	94.4287 Ops/s	94.8229 Ops/s	$\color{#d91a1a}-0.42\%$
test_vmap_transformer_speed[True-False]	10.0184ms	8.2659ms	120.9791 Ops/s	120.7509 Ops/s	$\color{#35bf28}+0.19\%$
test_vmap_transformer_speed[False-True]	23.5198ms	10.5735ms	94.5764 Ops/s	95.0052 Ops/s	$\color{#d91a1a}-0.45\%$
test_vmap_transformer_speed[False-False]	11.0682ms	8.2565ms	121.1172 Ops/s	122.2468 Ops/s	$\color{#d91a1a}-0.92\%$

[tcbegley]

LGTM!