[Doc] More doc on how to set and get non-tensor data

[vmoens]

No description provided.

[github-actions]

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

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

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_plain_set_nested	55.9940μs	16.2864μs	61.4011 KOps/s	62.7733 KOps/s	$\color{#d91a1a}-2.19\%$
test_plain_set_stack_nested	0.2462ms	0.1426ms	7.0136 KOps/s	6.9801 KOps/s	$\color{#35bf28}+0.48\%$
test_plain_set_nested_inplace	50.3750μs	18.6054μs	53.7477 KOps/s	54.8873 KOps/s	$\color{#d91a1a}-2.08\%$
test_plain_set_stack_nested_inplace	0.3074ms	0.1749ms	5.7192 KOps/s	5.6201 KOps/s	$\color{#35bf28}+1.76\%$
test_items	25.4380μs	2.4370μs	410.3444 KOps/s	412.1549 KOps/s	$\color{#d91a1a}-0.44\%$
test_items_nested	1.2223ms	0.2635ms	3.7956 KOps/s	3.5416 KOps/s	$\textbf{\color{#35bf28}+7.17\%}$
test_items_nested_locked	0.4029ms	0.2630ms	3.8017 KOps/s	3.5674 KOps/s	$\textbf{\color{#35bf28}+6.57\%}$
test_items_nested_leaf	0.6528ms	0.1630ms	6.1357 KOps/s	5.7346 KOps/s	$\textbf{\color{#35bf28}+6.99\%}$
test_items_stack_nested	1.5548ms	1.3011ms	768.6008 Ops/s	712.6116 Ops/s	$\textbf{\color{#35bf28}+7.86\%}$
test_items_stack_nested_leaf	1.6390ms	1.1715ms	853.6057 Ops/s	830.1804 Ops/s	$\color{#35bf28}+2.82\%$
test_items_stack_nested_locked	1.4602ms	0.8561ms	1.1680 KOps/s	1.1200 KOps/s	$\color{#35bf28}+4.29\%$
test_keys	17.9740μs	3.9340μs	254.1933 KOps/s	262.0829 KOps/s	$\color{#d91a1a}-3.01\%$
test_keys_nested	51.4457ms	0.1534ms	6.5190 KOps/s	6.7138 KOps/s	$\color{#d91a1a}-2.90\%$
test_keys_nested_locked	0.3049ms	0.1500ms	6.6687 KOps/s	6.6258 KOps/s	$\color{#35bf28}+0.65\%$
test_keys_nested_leaf	0.2815ms	0.1273ms	7.8534 KOps/s	7.7201 KOps/s	$\color{#35bf28}+1.73\%$
test_keys_stack_nested	2.0051ms	1.2523ms	798.5307 Ops/s	774.2813 Ops/s	$\color{#35bf28}+3.13\%$
test_keys_stack_nested_leaf	1.3641ms	1.2499ms	800.0630 Ops/s	772.4126 Ops/s	$\color{#35bf28}+3.58\%$
test_keys_stack_nested_locked	1.2998ms	0.7989ms	1.2517 KOps/s	1.2070 KOps/s	$\color{#35bf28}+3.70\%$
test_values	6.3242μs	1.1542μs	866.4318 KOps/s	878.2022 KOps/s	$\color{#d91a1a}-1.34\%$
test_values_nested	89.5180μs	50.3990μs	19.8417 KOps/s	19.4977 KOps/s	$\color{#35bf28}+1.76\%$
test_values_nested_locked	94.0960μs	50.8418μs	19.6688 KOps/s	19.4177 KOps/s	$\color{#35bf28}+1.29\%$
test_values_nested_leaf	91.6110μs	44.8562μs	22.2934 KOps/s	21.6083 KOps/s	$\color{#35bf28}+3.17\%$
test_values_stack_nested	1.2653ms	1.0210ms	979.4005 Ops/s	928.5735 Ops/s	$\textbf{\color{#35bf28}+5.47\%}$
test_values_stack_nested_leaf	1.3090ms	1.0182ms	982.0958 Ops/s	954.9510 Ops/s	$\color{#35bf28}+2.84\%$
test_values_stack_nested_locked	0.8452ms	0.6126ms	1.6323 KOps/s	1.6121 KOps/s	$\color{#35bf28}+1.25\%$
test_membership	20.5180μs	1.3290μs	752.4605 KOps/s	748.2442 KOps/s	$\color{#35bf28}+0.56\%$
test_membership_nested	20.0070μs	3.3671μs	296.9940 KOps/s	287.5979 KOps/s	$\color{#35bf28}+3.27\%$
test_membership_nested_leaf	27.7920μs	3.3826μs	295.6263 KOps/s	290.7656 KOps/s	$\color{#35bf28}+1.67\%$
test_membership_stacked_nested	35.0050μs	11.6798μs	85.6180 KOps/s	85.3908 KOps/s	$\color{#35bf28}+0.27\%$
test_membership_stacked_nested_leaf	32.0100μs	11.7154μs	85.3578 KOps/s	84.3367 KOps/s	$\color{#35bf28}+1.21\%$
test_membership_nested_last	27.7920μs	6.5508μs	152.6532 KOps/s	150.8709 KOps/s	$\color{#35bf28}+1.18\%$
test_membership_nested_leaf_last	27.9220μs	6.4650μs	154.6790 KOps/s	153.0863 KOps/s	$\color{#35bf28}+1.04\%$
test_membership_stacked_nested_last	0.3478ms	0.1745ms	5.7310 KOps/s	5.6624 KOps/s	$\color{#35bf28}+1.21\%$
test_membership_stacked_nested_leaf_last	39.2740μs	13.7718μs	72.6119 KOps/s	71.0375 KOps/s	$\color{#35bf28}+2.22\%$
test_nested_getleaf	39.7040μs	10.5148μs	95.1040 KOps/s	95.2130 KOps/s	$\color{#d91a1a}-0.11\%$
test_nested_get	39.6050μs	9.8090μs	101.9474 KOps/s	98.7760 KOps/s	$\color{#35bf28}+3.21\%$
test_stacked_getleaf	0.5805ms	0.3927ms	2.5462 KOps/s	2.5147 KOps/s	$\color{#35bf28}+1.25\%$
test_stacked_get	0.6457ms	0.3634ms	2.7520 KOps/s	2.7600 KOps/s	$\color{#d91a1a}-0.29\%$
test_nested_getitemleaf	44.2630μs	11.9049μs	83.9992 KOps/s	83.3901 KOps/s	$\color{#35bf28}+0.73\%$
test_nested_getitem	32.1300μs	11.2112μs	89.1962 KOps/s	87.6436 KOps/s	$\color{#35bf28}+1.77\%$
test_stacked_getitemleaf	0.6510ms	0.3951ms	2.5311 KOps/s	2.4491 KOps/s	$\color{#35bf28}+3.35\%$
test_stacked_getitem	0.4763ms	0.3653ms	2.7376 KOps/s	2.7149 KOps/s	$\color{#35bf28}+0.84\%$
test_lock_nested	0.7367ms	0.3290ms	3.0397 KOps/s	2.9449 KOps/s	$\color{#35bf28}+3.22\%$
test_lock_stack_nested	71.8615ms	5.4035ms	185.0664 Ops/s	178.1183 Ops/s	$\color{#35bf28}+3.90\%$
test_unlock_nested	0.7423ms	0.3332ms	3.0008 KOps/s	2.5135 KOps/s	$\textbf{\color{#35bf28}+19.39\%}$
test_unlock_stack_nested	84.3781ms	5.4677ms	182.8935 Ops/s	169.8776 Ops/s	$\textbf{\color{#35bf28}+7.66\%}$
test_flatten_speed	0.7940ms	0.3646ms	2.7426 KOps/s	2.7349 KOps/s	$\color{#35bf28}+0.28\%$
test_unflatten_speed	0.7880ms	0.4555ms	2.1956 KOps/s	2.1884 KOps/s	$\color{#35bf28}+0.33\%$
test_common_ops	4.3138ms	0.6730ms	1.4859 KOps/s	1.5181 KOps/s	$\color{#d91a1a}-2.12\%$
test_creation	60.4830μs	1.8597μs	537.7147 KOps/s	526.8642 KOps/s	$\color{#35bf28}+2.06\%$
test_creation_empty	28.8650μs	9.3544μs	106.9012 KOps/s	119.0314 KOps/s	$\textbf{\color{#d91a1a}-10.19\%}$
test_creation_nested_1	43.8920μs	12.0396μs	83.0590 KOps/s	90.1701 KOps/s	$\textbf{\color{#d91a1a}-7.89\%}$
test_creation_nested_2	37.8310μs	15.3597μs	65.1054 KOps/s	69.7347 KOps/s	$\textbf{\color{#d91a1a}-6.64\%}$
test_clone	86.5220μs	12.9932μs	76.9632 KOps/s	76.1193 KOps/s	$\color{#35bf28}+1.11\%$
test_getitem[int]	25.9090μs	10.9637μs	91.2101 KOps/s	89.9517 KOps/s	$\color{#35bf28}+1.40\%$
test_getitem[slice_int]	46.3370μs	21.8805μs	45.7028 KOps/s	43.3152 KOps/s	$\textbf{\color{#35bf28}+5.51\%}$
test_getitem[range]	0.1865ms	40.5120μs	24.6840 KOps/s	24.2052 KOps/s	$\color{#35bf28}+1.98\%$
test_getitem[tuple]	42.4390μs	17.7000μs	56.4971 KOps/s	55.2394 KOps/s	$\color{#35bf28}+2.28\%$
test_getitem[list]	80.4510μs	36.1263μs	27.6807 KOps/s	27.5743 KOps/s	$\color{#35bf28}+0.39\%$
test_setitem_dim[int]	46.3470μs	28.2310μs	35.4220 KOps/s	35.0032 KOps/s	$\color{#35bf28}+1.20\%$
test_setitem_dim[slice_int]	94.5560μs	55.0717μs	18.1582 KOps/s	17.9879 KOps/s	$\color{#35bf28}+0.95\%$
test_setitem_dim[range]	0.1181ms	72.7186μs	13.7516 KOps/s	13.8050 KOps/s	$\color{#d91a1a}-0.39\%$
test_setitem_dim[tuple]	61.9560μs	43.3224μs	23.0827 KOps/s	22.4587 KOps/s	$\color{#35bf28}+2.78\%$
test_setitem	88.4960μs	18.8029μs	53.1834 KOps/s	53.1418 KOps/s	$\color{#35bf28}+0.08\%$
test_set	80.6100μs	18.4450μs	54.2153 KOps/s	55.5018 KOps/s	$\color{#d91a1a}-2.32\%$
test_set_shared	3.4281ms	0.1431ms	6.9874 KOps/s	7.0740 KOps/s	$\color{#d91a1a}-1.22\%$
test_update	0.1128ms	20.7126μs	48.2798 KOps/s	50.2152 KOps/s	$\color{#d91a1a}-3.85\%$
test_update_nested	0.4480ms	29.0839μs	34.3833 KOps/s	35.6647 KOps/s	$\color{#d91a1a}-3.59\%$
test_set_nested	80.8310μs	20.1884μs	49.5334 KOps/s	50.0042 KOps/s	$\color{#d91a1a}-0.94\%$
test_set_nested_new	0.1218ms	23.6275μs	42.3236 KOps/s	41.0524 KOps/s	$\color{#35bf28}+3.10\%$
test_select	0.1181ms	36.4282μs	27.4513 KOps/s	26.5864 KOps/s	$\color{#35bf28}+3.25\%$
test_select_nested	0.1559ms	57.7298μs	17.3221 KOps/s	17.4680 KOps/s	$\color{#d91a1a}-0.84\%$
test_exclude_nested	0.2403ms	0.1077ms	9.2826 KOps/s	9.2041 KOps/s	$\color{#35bf28}+0.85\%$
test_empty[True]	0.4483ms	0.3215ms	3.1101 KOps/s	2.9651 KOps/s	$\color{#35bf28}+4.89\%$
test_empty[False]	30.1564μs	1.0361μs	965.1220 KOps/s	956.9889 KOps/s	$\color{#35bf28}+0.85\%$
test_unbind_speed	0.3003ms	0.2392ms	4.1810 KOps/s	3.9982 KOps/s	$\color{#35bf28}+4.57\%$
test_unbind_speed_stack0	72.2329ms	3.1985ms	312.6476 Ops/s	304.3285 Ops/s	$\color{#35bf28}+2.73\%$
test_unbind_speed_stack1	20.6980μs	1.9405μs	515.3363 KOps/s	506.9435 KOps/s	$\color{#35bf28}+1.66\%$
test_split	65.6168ms	1.6330ms	612.3822 Ops/s	600.0355 Ops/s	$\color{#35bf28}+2.06\%$
test_chunk	61.6766ms	1.5694ms	637.2031 Ops/s	629.4627 Ops/s	$\color{#35bf28}+1.23\%$
test_creation[device0]	3.6657ms	0.1032ms	9.6870 KOps/s	9.9777 KOps/s	$\color{#d91a1a}-2.91\%$
test_creation_from_tensor	0.1712ms	80.5729μs	12.4111 KOps/s	11.7025 KOps/s	$\textbf{\color{#35bf28}+6.05\%}$
test_add_one[memmap_tensor0]	0.2670ms	5.4822μs	182.4091 KOps/s	186.8885 KOps/s	$\color{#d91a1a}-2.40\%$
test_contiguous[memmap_tensor0]	11.8520μs	0.6538μs	1.5296 MOps/s	1.6213 MOps/s	$\textbf{\color{#d91a1a}-5.66\%}$
test_stack[memmap_tensor0]	59.6710μs	3.5241μs	283.7598 KOps/s	274.6673 KOps/s	$\color{#35bf28}+3.31\%$
test_memmaptd_index	1.0265ms	0.2189ms	4.5687 KOps/s	4.5704 KOps/s	$\color{#d91a1a}-0.04\%$
test_memmaptd_index_astensor	0.6624ms	0.2777ms	3.6004 KOps/s	3.5725 KOps/s	$\color{#35bf28}+0.78\%$
test_memmaptd_index_op	0.8715ms	0.5569ms	1.7957 KOps/s	1.8547 KOps/s	$\color{#d91a1a}-3.18\%$
test_serialize_model	0.1632s	0.1064s	9.3985 Ops/s	8.6583 Ops/s	$\textbf{\color{#35bf28}+8.55\%}$
test_serialize_model_pickle	0.4587s	0.3813s	2.6226 Ops/s	2.6015 Ops/s	$\color{#35bf28}+0.81\%$
test_serialize_weights	0.1668s	0.1070s	9.3437 Ops/s	9.0861 Ops/s	$\color{#35bf28}+2.83\%$
test_serialize_weights_returnearly	0.2890s	0.1449s	6.9028 Ops/s	7.6225 Ops/s	$\textbf{\color{#d91a1a}-9.44\%}$
test_serialize_weights_pickle	1.0197s	0.5833s	1.7144 Ops/s	2.4457 Ops/s	$\textbf{\color{#d91a1a}-29.90\%}$
test_serialize_weights_filesystem	0.1555s	98.7802ms	10.1235 Ops/s	10.9839 Ops/s	$\textbf{\color{#d91a1a}-7.83\%}$
test_serialize_model_filesystem	0.1026s	91.4042ms	10.9404 Ops/s	10.0135 Ops/s	$\textbf{\color{#35bf28}+9.26\%}$
test_reshape_pytree	60.1630μs	23.1584μs	43.1808 KOps/s	43.2826 KOps/s	$\color{#d91a1a}-0.24\%$
test_reshape_td	91.9910μs	29.7639μs	33.5977 KOps/s	33.2356 KOps/s	$\color{#35bf28}+1.09\%$
test_view_pytree	53.1490μs	22.8999μs	43.6682 KOps/s	43.4180 KOps/s	$\color{#35bf28}+0.58\%$
test_view_td	25.2770μs	4.9700μs	201.2062 KOps/s	204.7399 KOps/s	$\color{#d91a1a}-1.73\%$
test_unbind_pytree	92.4920μs	26.2713μs	38.0643 KOps/s	37.9929 KOps/s	$\color{#35bf28}+0.19\%$
test_unbind_td	0.4668ms	35.8940μs	27.8598 KOps/s	27.8363 KOps/s	$\color{#35bf28}+0.08\%$
test_split_pytree	75.3610μs	26.5035μs	37.7309 KOps/s	37.9556 KOps/s	$\color{#d91a1a}-0.59\%$
test_split_td	0.1135ms	39.8915μs	25.0680 KOps/s	24.6941 KOps/s	$\color{#35bf28}+1.51\%$
test_add_pytree	73.0560μs	31.9853μs	31.2643 KOps/s	31.4438 KOps/s	$\color{#d91a1a}-0.57\%$
test_add_td	0.1152ms	48.9540μs	20.4273 KOps/s	20.9668 KOps/s	$\color{#d91a1a}-2.57\%$
test_distributed	0.4907ms	96.9925μs	10.3101 KOps/s	9.9001 KOps/s	$\color{#35bf28}+4.14\%$
test_tdmodule	0.1044ms	21.6292μs	46.2338 KOps/s	46.1685 KOps/s	$\color{#35bf28}+0.14\%$
test_tdmodule_dispatch	0.2512ms	42.7931μs	23.3683 KOps/s	24.0203 KOps/s	$\color{#d91a1a}-2.71\%$
test_tdseq	0.3741ms	25.6190μs	39.0335 KOps/s	39.7671 KOps/s	$\color{#d91a1a}-1.84\%$
test_tdseq_dispatch	0.1527ms	47.4863μs	21.0587 KOps/s	21.9301 KOps/s	$\color{#d91a1a}-3.97\%$
test_instantiation_functorch	1.6002ms	1.2820ms	780.0263 Ops/s	755.5419 Ops/s	$\color{#35bf28}+3.24\%$
test_instantiation_td	1.4548ms	0.9968ms	1.0032 KOps/s	982.9805 Ops/s	$\color{#35bf28}+2.05\%$
test_exec_functorch	0.2981ms	0.1638ms	6.1037 KOps/s	6.2160 KOps/s	$\color{#d91a1a}-1.81\%$
test_exec_functional_call	0.3128ms	0.1492ms	6.7036 KOps/s	6.7783 KOps/s	$\color{#d91a1a}-1.10\%$
test_exec_td	0.2884ms	0.1468ms	6.8133 KOps/s	7.0375 KOps/s	$\color{#d91a1a}-3.19\%$
test_exec_td_decorator	0.9075ms	0.1821ms	5.4922 KOps/s	5.6136 KOps/s	$\color{#d91a1a}-2.16\%$
test_vmap_mlp_speed[True-True]	1.4445ms	0.8753ms	1.1425 KOps/s	1.1244 KOps/s	$\color{#35bf28}+1.60\%$
test_vmap_mlp_speed[True-False]	0.6369ms	0.4622ms	2.1633 KOps/s	2.1599 KOps/s	$\color{#35bf28}+0.16\%$
test_vmap_mlp_speed[False-True]	1.1181ms	0.7577ms	1.3198 KOps/s	1.2861 KOps/s	$\color{#35bf28}+2.62\%$
test_vmap_mlp_speed[False-False]	0.4753ms	0.3813ms	2.6226 KOps/s	2.5911 KOps/s	$\color{#35bf28}+1.21\%$
test_vmap_mlp_speed_decorator[True-True]	2.8521ms	2.2424ms	445.9589 Ops/s	438.5254 Ops/s	$\color{#35bf28}+1.70\%$
test_vmap_mlp_speed_decorator[True-False]	0.8685ms	0.5097ms	1.9620 KOps/s	1.9257 KOps/s	$\color{#35bf28}+1.88\%$
test_vmap_mlp_speed_decorator[False-True]	3.6388ms	1.8344ms	545.1362 Ops/s	536.3678 Ops/s	$\color{#35bf28}+1.63\%$
test_vmap_mlp_speed_decorator[False-False]	0.7580ms	0.3965ms	2.5218 KOps/s	2.3423 KOps/s	$\textbf{\color{#35bf28}+7.66\%}$

[github-actions]

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

Total Benchmarks: 132. Improved: $\large\color{#35bf28}7$. Worsened: $\large\color{#d91a1a}29$.

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_plain_set_nested	0.1350ms	14.6658μs	68.1858 KOps/s	74.7710 KOps/s	$\textbf{\color{#d91a1a}-8.81\%}$
test_plain_set_stack_nested	0.1432ms	0.1214ms	8.2359 KOps/s	8.3583 KOps/s	$\color{#d91a1a}-1.46\%$
test_plain_set_nested_inplace	46.5910μs	16.1496μs	61.9210 KOps/s	67.8877 KOps/s	$\textbf{\color{#d91a1a}-8.79\%}$
test_plain_set_stack_nested_inplace	0.2144ms	0.1509ms	6.6263 KOps/s	6.7588 KOps/s	$\color{#d91a1a}-1.96\%$
test_items	27.3910μs	4.7840μs	209.0315 KOps/s	208.7994 KOps/s	$\color{#35bf28}+0.11\%$
test_items_nested	0.3682ms	0.3392ms	2.9483 KOps/s	2.9281 KOps/s	$\color{#35bf28}+0.69\%$
test_items_nested_locked	0.3790ms	0.3438ms	2.9089 KOps/s	2.8824 KOps/s	$\color{#35bf28}+0.92\%$
test_items_nested_leaf	0.2258ms	0.1997ms	5.0082 KOps/s	4.9434 KOps/s	$\color{#35bf28}+1.31\%$
test_items_stack_nested	1.3559ms	1.3113ms	762.6157 Ops/s	758.5045 Ops/s	$\color{#35bf28}+0.54\%$
test_items_stack_nested_leaf	1.2082ms	1.1564ms	864.7537 Ops/s	867.1391 Ops/s	$\color{#d91a1a}-0.28\%$
test_items_stack_nested_locked	0.9419ms	0.8956ms	1.1165 KOps/s	1.0987 KOps/s	$\color{#35bf28}+1.62\%$
test_keys	24.8910μs	4.6215μs	216.3784 KOps/s	214.7416 KOps/s	$\color{#35bf28}+0.76\%$
test_keys_nested	1.7406ms	95.7756μs	10.4411 KOps/s	10.4691 KOps/s	$\color{#d91a1a}-0.27\%$
test_keys_nested_locked	0.1309ms	98.8766μs	10.1136 KOps/s	10.1056 KOps/s	$\color{#35bf28}+0.08\%$
test_keys_nested_leaf	0.1824ms	78.9085μs	12.6729 KOps/s	12.6618 KOps/s	$\color{#35bf28}+0.09\%$
test_keys_stack_nested	1.2136ms	1.1646ms	858.6592 Ops/s	879.6882 Ops/s	$\color{#d91a1a}-2.39\%$
test_keys_stack_nested_leaf	1.1734ms	1.1443ms	873.9175 Ops/s	879.5197 Ops/s	$\color{#d91a1a}-0.64\%$
test_keys_stack_nested_locked	0.7840ms	0.7214ms	1.3862 KOps/s	1.3666 KOps/s	$\color{#35bf28}+1.44\%$
test_values	6.5800μs	1.8940μs	527.9853 KOps/s	525.0030 KOps/s	$\color{#35bf28}+0.57\%$
test_values_nested	70.8210μs	45.5713μs	21.9437 KOps/s	21.9142 KOps/s	$\color{#35bf28}+0.13\%$
test_values_nested_locked	73.2310μs	48.1822μs	20.7546 KOps/s	20.8784 KOps/s	$\color{#d91a1a}-0.59\%$
test_values_nested_leaf	59.4610μs	39.6437μs	25.2247 KOps/s	25.0536 KOps/s	$\color{#35bf28}+0.68\%$
test_values_stack_nested	1.0065ms	0.9638ms	1.0376 KOps/s	1.0368 KOps/s	$\color{#35bf28}+0.08\%$
test_values_stack_nested_leaf	1.0334ms	0.9548ms	1.0473 KOps/s	1.0435 KOps/s	$\color{#35bf28}+0.36\%$
test_values_stack_nested_locked	0.6438ms	0.5803ms	1.7233 KOps/s	1.7508 KOps/s	$\color{#d91a1a}-1.57\%$
test_membership	4.9900μs	0.9692μs	1.0318 MOps/s	919.8546 KOps/s	$\textbf{\color{#35bf28}+12.17\%}$
test_membership_nested	29.5510μs	2.9062μs	344.0915 KOps/s	338.5157 KOps/s	$\color{#35bf28}+1.65\%$
test_membership_nested_leaf	29.1510μs	2.9093μs	343.7287 KOps/s	337.1853 KOps/s	$\color{#35bf28}+1.94\%$
test_membership_stacked_nested	30.5110μs	11.4915μs	87.0206 KOps/s	87.7584 KOps/s	$\color{#d91a1a}-0.84\%$
test_membership_stacked_nested_leaf	31.0210μs	11.7245μs	85.2911 KOps/s	87.9540 KOps/s	$\color{#d91a1a}-3.03\%$
test_membership_nested_last	43.7910μs	5.3970μs	185.2880 KOps/s	186.3074 KOps/s	$\color{#d91a1a}-0.55\%$
test_membership_nested_leaf_last	35.6210μs	5.3647μs	186.4049 KOps/s	186.1145 KOps/s	$\color{#35bf28}+0.16\%$
test_membership_stacked_nested_last	0.2003ms	0.1578ms	6.3369 KOps/s	6.4344 KOps/s	$\color{#d91a1a}-1.51\%$
test_membership_stacked_nested_leaf_last	29.4500μs	13.4421μs	74.3931 KOps/s	76.4191 KOps/s	$\color{#d91a1a}-2.65\%$
test_nested_getleaf	32.0710μs	8.5414μs	117.0774 KOps/s	119.1763 KOps/s	$\color{#d91a1a}-1.76\%$
test_nested_get	39.0810μs	8.0436μs	124.3220 KOps/s	126.6829 KOps/s	$\color{#d91a1a}-1.86\%$
test_stacked_getleaf	0.3708ms	0.3337ms	2.9970 KOps/s	3.0369 KOps/s	$\color{#d91a1a}-1.31\%$
test_stacked_get	0.3388ms	0.3017ms	3.3148 KOps/s	3.3534 KOps/s	$\color{#d91a1a}-1.15\%$
test_nested_getitemleaf	32.0800μs	9.9450μs	100.5532 KOps/s	102.0204 KOps/s	$\color{#d91a1a}-1.44\%$
test_nested_getitem	35.1010μs	9.4746μs	105.5448 KOps/s	106.9335 KOps/s	$\color{#d91a1a}-1.30\%$
test_stacked_getitemleaf	0.3899ms	0.3343ms	2.9909 KOps/s	2.9967 KOps/s	$\color{#d91a1a}-0.19\%$
test_stacked_getitem	0.3723ms	0.3011ms	3.3209 KOps/s	3.2971 KOps/s	$\color{#35bf28}+0.72\%$
test_lock_nested	6.5062ms	0.3582ms	2.7919 KOps/s	2.8679 KOps/s	$\color{#d91a1a}-2.65\%$
test_lock_stack_nested	83.9937ms	6.2716ms	159.4498 Ops/s	159.0681 Ops/s	$\color{#35bf28}+0.24\%$
test_unlock_nested	75.5661ms	0.4223ms	2.3682 KOps/s	2.9009 KOps/s	$\textbf{\color{#d91a1a}-18.36\%}$
test_unlock_stack_nested	83.5604ms	6.2773ms	159.3036 Ops/s	159.2724 Ops/s	$\color{#35bf28}+0.02\%$
test_flatten_speed	0.5148ms	0.2626ms	3.8083 KOps/s	3.7840 KOps/s	$\color{#35bf28}+0.64\%$
test_unflatten_speed	0.3984ms	0.3646ms	2.7426 KOps/s	2.7725 KOps/s	$\color{#d91a1a}-1.08\%$
test_common_ops	1.1853ms	0.6340ms	1.5774 KOps/s	1.5330 KOps/s	$\color{#35bf28}+2.89\%$
test_creation	16.5310μs	1.5845μs	631.1217 KOps/s	635.5466 KOps/s	$\color{#d91a1a}-0.70\%$
test_creation_empty	33.0010μs	10.2214μs	97.8341 KOps/s	126.6855 KOps/s	$\textbf{\color{#d91a1a}-22.77\%}$
test_creation_nested_1	41.6810μs	11.9824μs	83.4557 KOps/s	104.9738 KOps/s	$\textbf{\color{#d91a1a}-20.50\%}$
test_creation_nested_2	34.3010μs	14.3671μs	69.6033 KOps/s	82.6072 KOps/s	$\textbf{\color{#d91a1a}-15.74\%}$
test_clone	0.1071ms	13.8155μs	72.3825 KOps/s	72.8836 KOps/s	$\color{#d91a1a}-0.69\%$
test_getitem[int]	25.4900μs	10.9305μs	91.4872 KOps/s	94.9051 KOps/s	$\color{#d91a1a}-3.60\%$
test_getitem[slice_int]	51.2810μs	21.5255μs	46.4566 KOps/s	49.2860 KOps/s	$\textbf{\color{#d91a1a}-5.74\%}$
test_getitem[range]	69.1710μs	38.0394μs	26.2885 KOps/s	27.5508 KOps/s	$\color{#d91a1a}-4.58\%$
test_getitem[tuple]	40.0510μs	18.9914μs	52.6555 KOps/s	53.7151 KOps/s	$\color{#d91a1a}-1.97\%$
test_getitem[list]	0.3246ms	33.4875μs	29.8619 KOps/s	31.0010 KOps/s	$\color{#d91a1a}-3.67\%$
test_setitem_dim[int]	47.7510μs	28.2533μs	35.3941 KOps/s	38.5427 KOps/s	$\textbf{\color{#d91a1a}-8.17\%}$
test_setitem_dim[slice_int]	69.7510μs	49.8087μs	20.0768 KOps/s	20.7763 KOps/s	$\color{#d91a1a}-3.37\%$
test_setitem_dim[range]	89.3720μs	63.9323μs	15.6416 KOps/s	16.4345 KOps/s	$\color{#d91a1a}-4.83\%$
test_setitem_dim[tuple]	71.7010μs	43.3779μs	23.0532 KOps/s	24.6624 KOps/s	$\textbf{\color{#d91a1a}-6.53\%}$
test_setitem	0.1127ms	19.5100μs	51.2558 KOps/s	55.0047 KOps/s	$\textbf{\color{#d91a1a}-6.82\%}$
test_set	0.1091ms	19.0895μs	52.3847 KOps/s	56.8390 KOps/s	$\textbf{\color{#d91a1a}-7.84\%}$
test_set_shared	2.9984ms	0.1072ms	9.3257 KOps/s	9.7753 KOps/s	$\color{#d91a1a}-4.60\%$
test_update	0.1134ms	22.6417μs	44.1664 KOps/s	50.8521 KOps/s	$\textbf{\color{#d91a1a}-13.15\%}$
test_update_nested	0.1093ms	29.0195μs	34.4596 KOps/s	38.6542 KOps/s	$\textbf{\color{#d91a1a}-10.85\%}$
test_set_nested	0.1166ms	20.5660μs	48.6240 KOps/s	54.3460 KOps/s	$\textbf{\color{#d91a1a}-10.53\%}$
test_set_nested_new	0.1104ms	23.2151μs	43.0754 KOps/s	47.0148 KOps/s	$\textbf{\color{#d91a1a}-8.38\%}$
test_select	0.1099ms	36.8006μs	27.1735 KOps/s	30.3424 KOps/s	$\textbf{\color{#d91a1a}-10.44\%}$
test_select_nested	81.3020μs	54.1217μs	18.4769 KOps/s	18.8743 KOps/s	$\color{#d91a1a}-2.11\%$
test_exclude_nested	0.1343ms	0.1075ms	9.2991 KOps/s	9.3574 KOps/s	$\color{#d91a1a}-0.62\%$
test_empty[True]	0.3556ms	0.3182ms	3.1430 KOps/s	3.0864 KOps/s	$\color{#35bf28}+1.83\%$
test_empty[False]	2.8110μs	0.8655μs	1.1554 MOps/s	1.1621 MOps/s	$\color{#d91a1a}-0.58\%$
test_to	73.9210μs	54.3592μs	18.3962 KOps/s	19.1868 KOps/s	$\color{#d91a1a}-4.12\%$
test_to_nonblocking	56.4330μs	34.6725μs	28.8413 KOps/s	30.6528 KOps/s	$\textbf{\color{#d91a1a}-5.91\%}$
test_unbind_speed	0.2939ms	0.2677ms	3.7359 KOps/s	3.8251 KOps/s	$\color{#d91a1a}-2.33\%$
test_unbind_speed_stack0	81.3480ms	3.6904ms	270.9737 Ops/s	264.9235 Ops/s	$\color{#35bf28}+2.28\%$
test_unbind_speed_stack1	8.7100μs	1.7356μs	576.1705 KOps/s	584.6721 KOps/s	$\color{#d91a1a}-1.45\%$
test_split	76.3250ms	1.7226ms	580.5063 Ops/s	657.1994 Ops/s	$\textbf{\color{#d91a1a}-11.67\%}$
test_chunk	1.5925ms	1.5405ms	649.1414 Ops/s	610.2404 Ops/s	$\textbf{\color{#35bf28}+6.37\%}$
test_creation[device0]	0.1350ms	73.2239μs	13.6567 KOps/s	13.7694 KOps/s	$\color{#d91a1a}-0.82\%$
test_creation_from_tensor	0.1271ms	53.6139μs	18.6519 KOps/s	18.6173 KOps/s	$\color{#35bf28}+0.19\%$
test_add_one[memmap_tensor0]	0.2507ms	6.9440μs	144.0094 KOps/s	144.6491 KOps/s	$\color{#d91a1a}-0.44\%$
test_contiguous[memmap_tensor0]	26.2100μs	0.6356μs	1.5732 MOps/s	1.5320 MOps/s	$\color{#35bf28}+2.69\%$
test_stack[memmap_tensor0]	53.6220μs	4.5940μs	217.6755 KOps/s	219.5027 KOps/s	$\color{#d91a1a}-0.83\%$
test_memmaptd_index	77.4281ms	0.2940ms	3.4018 KOps/s	3.9756 KOps/s	$\textbf{\color{#d91a1a}-14.43\%}$
test_memmaptd_index_astensor	0.6101ms	0.3176ms	3.1489 KOps/s	3.2059 KOps/s	$\color{#d91a1a}-1.78\%$
test_memmaptd_index_op	0.9856ms	0.6411ms	1.5598 KOps/s	1.7082 KOps/s	$\textbf{\color{#d91a1a}-8.69\%}$
test_serialize_model	0.1693s	97.5555ms	10.2506 Ops/s	9.5970 Ops/s	$\textbf{\color{#35bf28}+6.81\%}$
test_serialize_model_pickle	1.3464s	1.2355s	0.8094 Ops/s	0.8070 Ops/s	$\color{#35bf28}+0.29\%$
test_serialize_weights	88.6717ms	86.6459ms	11.5412 Ops/s	9.9591 Ops/s	$\textbf{\color{#35bf28}+15.89\%}$
test_serialize_weights_returnearly	0.2400s	72.3579ms	13.8202 Ops/s	12.6060 Ops/s	$\textbf{\color{#35bf28}+9.63\%}$
test_serialize_weights_pickle	1.3554s	1.2367s	0.8086 Ops/s	0.8080 Ops/s	$\color{#35bf28}+0.08\%$
test_reshape_pytree	0.2467ms	25.6155μs	39.0389 KOps/s	39.1997 KOps/s	$\color{#d91a1a}-0.41\%$
test_reshape_td	0.1697ms	31.5528μs	31.6929 KOps/s	33.8228 KOps/s	$\textbf{\color{#d91a1a}-6.30\%}$
test_view_pytree	54.0520μs	25.8757μs	38.6463 KOps/s	39.8355 KOps/s	$\color{#d91a1a}-2.99\%$
test_view_td	20.6510μs	4.2031μs	237.9172 KOps/s	235.6812 KOps/s	$\color{#35bf28}+0.95\%$
test_unbind_pytree	46.7210μs	30.0355μs	33.2939 KOps/s	33.7908 KOps/s	$\color{#d91a1a}-1.47\%$
test_unbind_td	0.5208ms	40.4090μs	24.7470 KOps/s	25.0974 KOps/s	$\color{#d91a1a}-1.40\%$
test_split_pytree	67.2510μs	29.2858μs	34.1462 KOps/s	34.1465 KOps/s	$-0.00\%$
test_split_td	0.1237ms	38.7966μs	25.7754 KOps/s	25.8940 KOps/s	$\color{#d91a1a}-0.46\%$
test_add_pytree	63.0510μs	35.2915μs	28.3354 KOps/s	26.6685 KOps/s	$\textbf{\color{#35bf28}+6.25\%}$
test_add_td	0.1302ms	50.2903μs	19.8846 KOps/s	21.6401 KOps/s	$\textbf{\color{#d91a1a}-8.11\%}$
test_distributed	0.1813ms	69.4402μs	14.4009 KOps/s	13.6848 KOps/s	$\textbf{\color{#35bf28}+5.23\%}$
test_tdmodule	0.1123ms	18.7594μs	53.3065 KOps/s	57.7998 KOps/s	$\textbf{\color{#d91a1a}-7.77\%}$
test_tdmodule_dispatch	0.2211ms	38.9917μs	25.6465 KOps/s	27.8855 KOps/s	$\textbf{\color{#d91a1a}-8.03\%}$
test_tdseq	35.8710μs	21.5816μs	46.3357 KOps/s	49.1130 KOps/s	$\textbf{\color{#d91a1a}-5.65\%}$
test_tdseq_dispatch	60.5220μs	41.6522μs	24.0083 KOps/s	25.8752 KOps/s	$\textbf{\color{#d91a1a}-7.21\%}$
test_instantiation_functorch	1.7477ms	1.6494ms	606.2655 Ops/s	601.0588 Ops/s	$\color{#35bf28}+0.87\%$
test_instantiation_td	1.6935ms	1.1484ms	870.7695 Ops/s	852.4701 Ops/s	$\color{#35bf28}+2.15\%$
test_exec_functorch	0.2360ms	0.1595ms	6.2706 KOps/s	6.1960 KOps/s	$\color{#35bf28}+1.20\%$
test_exec_functional_call	0.2253ms	0.1592ms	6.2795 KOps/s	6.1750 KOps/s	$\color{#35bf28}+1.69\%$
test_exec_td	0.1782ms	0.1496ms	6.6833 KOps/s	6.7354 KOps/s	$\color{#d91a1a}-0.77\%$
test_exec_td_decorator	0.7471ms	0.1897ms	5.2721 KOps/s	5.2644 KOps/s	$\color{#35bf28}+0.15\%$
test_vmap_mlp_speed[True-True]	1.4633ms	1.1007ms	908.5459 Ops/s	917.0771 Ops/s	$\color{#d91a1a}-0.93\%$
test_vmap_mlp_speed[True-False]	0.8100ms	0.6620ms	1.5105 KOps/s	1.5050 KOps/s	$\color{#35bf28}+0.37\%$
test_vmap_mlp_speed[False-True]	1.1037ms	1.0057ms	994.3741 Ops/s	952.8571 Ops/s	$\color{#35bf28}+4.36\%$
test_vmap_mlp_speed[False-False]	0.6416ms	0.5755ms	1.7375 KOps/s	1.6730 KOps/s	$\color{#35bf28}+3.85\%$
test_vmap_mlp_speed_decorator[True-True]	2.9465ms	2.4360ms	410.5127 Ops/s	419.7196 Ops/s	$\color{#d91a1a}-2.19\%$
test_vmap_mlp_speed_decorator[True-False]	1.0391ms	0.6915ms	1.4462 KOps/s	1.3915 KOps/s	$\color{#35bf28}+3.94\%$
test_vmap_mlp_speed_decorator[False-True]	2.4385ms	2.0613ms	485.1417 Ops/s	497.2363 Ops/s	$\color{#d91a1a}-2.43\%$
test_vmap_mlp_speed_decorator[False-False]	1.0857ms	0.6275ms	1.5937 KOps/s	1.6779 KOps/s	$\textbf{\color{#d91a1a}-5.02\%}$
test_vmap_transformer_speed[True-True]	13.0559ms	12.6069ms	79.3216 Ops/s	81.2635 Ops/s	$\color{#d91a1a}-2.39\%$
test_vmap_transformer_speed[True-False]	8.8354ms	8.2635ms	121.0146 Ops/s	124.0553 Ops/s	$\color{#d91a1a}-2.45\%$
test_vmap_transformer_speed[False-True]	13.0123ms	12.5092ms	79.9411 Ops/s	81.8017 Ops/s	$\color{#d91a1a}-2.27\%$
test_vmap_transformer_speed[False-False]	8.7825ms	8.3035ms	120.4308 Ops/s	125.0926 Ops/s	$\color{#d91a1a}-3.73\%$
test_vmap_transformer_speed_decorator[True-True]	0.1710s	82.6110ms	12.1049 Ops/s	13.7116 Ops/s	$\textbf{\color{#d91a1a}-11.72\%}$
test_vmap_transformer_speed_decorator[True-False]	21.7566ms	19.7891ms	50.5330 Ops/s	51.3548 Ops/s	$\color{#d91a1a}-1.60\%$
test_vmap_transformer_speed_decorator[False-True]	69.4447ms	68.2609ms	14.6497 Ops/s	15.2936 Ops/s	$\color{#d91a1a}-4.21\%$
test_vmap_transformer_speed_decorator[False-False]	0.1191s	21.2319ms	47.0990 Ops/s	53.0165 Ops/s	$\textbf{\color{#d91a1a}-11.16\%}$

[shagunsodhani]

curious would this still work if the batch_size was 2 dimensional (say (3, 2)). Similarly, would

assert td[0][1]["b"] == "a string!"

work with a 2-dimensional batch size ?