[Feature] Unravel nested keys

[vmoens]

Description

cc @Blonck

[github-actions]

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

Total Benchmarks: 47. Improved: $\large\color{#35bf28}10$. Worsened: $\large\color{#d91a1a}10$.

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_common_ops	1.9576ms	1.7186ms	581.8604 Ops/s	621.2938 Ops/s	$\textbf{\color{#d91a1a}-6.35\%}$
test_creation	12.7421μs	4.4451μs	224.9693 KOps/s	282.5267 KOps/s	$\textbf{\color{#d91a1a}-20.37\%}$
test_creation_empty	60.2533μs	18.4167μs	54.2985 KOps/s	66.2349 KOps/s	$\textbf{\color{#d91a1a}-18.02\%}$
test_creation_nested_1	47.2102μs	30.5785μs	32.7028 KOps/s	38.2706 KOps/s	$\textbf{\color{#d91a1a}-14.55\%}$
test_creation_nested_2	62.0233μs	31.1056μs	32.1486 KOps/s	34.6065 KOps/s	$\textbf{\color{#d91a1a}-7.10\%}$
test_clone	40.7892μs	30.8219μs	32.4444 KOps/s	34.6741 KOps/s	$\textbf{\color{#d91a1a}-6.43\%}$
test_getitem[int]	62.2441μs	38.7814μs	25.7856 KOps/s	27.3651 KOps/s	$\textbf{\color{#d91a1a}-5.77\%}$
test_getitem[slice_int]	0.1204ms	82.6325μs	12.1018 KOps/s	12.0938 KOps/s	$\color{#35bf28}+0.07\%$
test_getitem[range]	0.1538ms	0.1137ms	8.7974 KOps/s	9.4798 KOps/s	$\textbf{\color{#d91a1a}-7.20\%}$
test_getitem[tuple]	0.1084ms	69.6457μs	14.3584 KOps/s	13.9544 KOps/s	$\color{#35bf28}+2.90\%$
test_getitem[list]	0.1283ms	95.3376μs	10.4890 KOps/s	11.1472 KOps/s	$\textbf{\color{#d91a1a}-5.90\%}$
test_setitem_dim[int]	3.0604ms	62.6996μs	15.9491 KOps/s	15.3116 KOps/s	$\color{#35bf28}+4.16\%$
test_setitem_dim[slice_int]	3.1551ms	0.1144ms	8.7447 KOps/s	8.4828 KOps/s	$\color{#35bf28}+3.09\%$
test_setitem_dim[range]	4.1746ms	0.1276ms	7.8357 KOps/s	8.0284 KOps/s	$\color{#d91a1a}-2.40\%$
test_setitem_dim[tuple]	3.2840ms	0.1011ms	9.8927 KOps/s	10.1506 KOps/s	$\color{#d91a1a}-2.54\%$
test_setitem	75.1254μs	46.8675μs	21.3368 KOps/s	22.2954 KOps/s	$\color{#d91a1a}-4.30\%$
test_set	70.9723μs	45.0705μs	22.1874 KOps/s	22.0262 KOps/s	$\color{#35bf28}+0.73\%$
test_set_shared	0.4315ms	0.3184ms	3.1405 KOps/s	3.1166 KOps/s	$\color{#35bf28}+0.77\%$
test_update	0.1175ms	57.4980μs	17.3919 KOps/s	16.6127 KOps/s	$\color{#35bf28}+4.69\%$
test_update_nested	0.1123ms	79.3652μs	12.6000 KOps/s	11.8829 KOps/s	$\textbf{\color{#35bf28}+6.03\%}$
test_set_nested	0.1180ms	56.6622μs	17.6484 KOps/s	18.2952 KOps/s	$\color{#d91a1a}-3.54\%$
test_set_nested_new	0.1069ms	74.7894μs	13.3709 KOps/s	13.8365 KOps/s	$\color{#d91a1a}-3.37\%$
test_select	0.1506ms	0.1189ms	8.4099 KOps/s	8.6348 KOps/s	$\color{#d91a1a}-2.60\%$
test_creation[device0]	1.4779ms	0.5850ms	1.7095 KOps/s	1.7884 KOps/s	$\color{#d91a1a}-4.41\%$
test_creation_from_tensor	0.7404ms	0.5902ms	1.6942 KOps/s	1.8661 KOps/s	$\textbf{\color{#d91a1a}-9.21\%}$
test_add_one[memmap_tensor0]	67.2553μs	55.5963μs	17.9868 KOps/s	17.6947 KOps/s	$\color{#35bf28}+1.65\%$
test_contiguous[memmap_tensor0]	13.3401μs	10.1866μs	98.1678 KOps/s	86.6302 KOps/s	$\textbf{\color{#35bf28}+13.32\%}$
test_stack[memmap_tensor0]	0.1915ms	59.0503μs	16.9347 KOps/s	6.6551 KOps/s	$\textbf{\color{#35bf28}+154.46\%}$
test_reshape_pytree	44.3262μs	34.1205μs	29.3079 KOps/s	27.3791 KOps/s	$\textbf{\color{#35bf28}+7.04\%}$
test_reshape_td	0.1585ms	52.9433μs	18.8881 KOps/s	18.8900 KOps/s	$-0.01\%$
test_view_pytree	67.7453μs	31.8269μs	31.4200 KOps/s	30.4477 KOps/s	$\color{#35bf28}+3.19\%$
test_view_td	11.5461μs	7.8454μs	127.4636 KOps/s	113.5320 KOps/s	$\textbf{\color{#35bf28}+12.27\%}$
test_unbind_pytree	45.8192μs	35.8179μs	27.9190 KOps/s	25.6732 KOps/s	$\textbf{\color{#35bf28}+8.75\%}$
test_unbind_td	0.2590ms	0.2144ms	4.6644 KOps/s	4.6598 KOps/s	$\color{#35bf28}+0.10\%$
test_split_pytree	55.4763μs	41.9286μs	23.8500 KOps/s	22.1070 KOps/s	$\textbf{\color{#35bf28}+7.88\%}$
test_split_td	0.1906ms	0.1413ms	7.0751 KOps/s	7.2484 KOps/s	$\color{#d91a1a}-2.39\%$
test_add_pytree	63.2703μs	52.9465μs	18.8870 KOps/s	18.8684 KOps/s	$\color{#35bf28}+0.10\%$
test_add_td	0.1565ms	0.1140ms	8.7719 KOps/s	9.1582 KOps/s	$\color{#d91a1a}-4.22\%$
test_distributed	0.1082ms	0.1082ms	9.2421 KOps/s	8.8339 KOps/s	$\color{#35bf28}+4.62\%$
test_tdmodule	3.1514ms	38.6048μs	25.9035 KOps/s	27.1761 KOps/s	$\color{#d91a1a}-4.68\%$
test_tdmodule_dispatch	6.5432ms	81.6184μs	12.2521 KOps/s	11.5535 KOps/s	$\textbf{\color{#35bf28}+6.05\%}$
test_tdseq	1.9082ms	48.1872μs	20.7524 KOps/s	20.8072 KOps/s	$\color{#d91a1a}-0.26\%$
test_tdseq_dispatch	3.9636ms	98.1804μs	10.1853 KOps/s	10.3515 KOps/s	$\color{#d91a1a}-1.61\%$
test_instantiation_functorch	2.0978ms	1.6828ms	594.2377 Ops/s	542.7067 Ops/s	$\textbf{\color{#35bf28}+9.50\%}$
test_instantiation_td	2.0707ms	1.3046ms	766.5197 Ops/s	737.2330 Ops/s	$\color{#35bf28}+3.97\%$
test_exec_functorch	0.3422ms	0.2628ms	3.8053 KOps/s	3.4786 KOps/s	$\textbf{\color{#35bf28}+9.39\%}$
test_exec_td	0.5223ms	0.4636ms	2.1571 KOps/s	2.1556 KOps/s	$\color{#35bf28}+0.07\%$

[github-actions]

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

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

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_common_ops	2.4356ms	1.9403ms	515.3787 Ops/s	517.4126 Ops/s	$\color{#d91a1a}-0.39\%$
test_creation	9.0103μs	5.2194μs	191.5939 KOps/s	196.6904 KOps/s	$\color{#d91a1a}-2.59\%$
test_creation_empty	26.7249μs	20.9649μs	47.6988 KOps/s	47.4778 KOps/s	$\color{#35bf28}+0.47\%$
test_creation_nested_1	67.6812μs	38.1857μs	26.1878 KOps/s	26.5816 KOps/s	$\color{#d91a1a}-1.48\%$
test_creation_nested_2	97.4012μs	41.2235μs	24.2580 KOps/s	25.2092 KOps/s	$\color{#d91a1a}-3.77\%$
test_clone	52.0587μs	39.0810μs	25.5879 KOps/s	25.9853 KOps/s	$\color{#d91a1a}-1.53\%$
test_getitem[int]	87.7947μs	46.4177μs	21.5435 KOps/s	21.5365 KOps/s	$\color{#35bf28}+0.03\%$
test_getitem[slice_int]	0.1821ms	0.1041ms	9.6100 KOps/s	9.4885 KOps/s	$\color{#35bf28}+1.28\%$
test_getitem[range]	0.1694ms	0.1279ms	7.8175 KOps/s	7.8088 KOps/s	$\color{#35bf28}+0.11\%$
test_getitem[tuple]	0.1318ms	90.5834μs	11.0396 KOps/s	11.0023 KOps/s	$\color{#35bf28}+0.34\%$
test_getitem[list]	0.1396ms	0.1088ms	9.1934 KOps/s	9.1215 KOps/s	$\color{#35bf28}+0.79\%$
test_setitem_dim[int]	3.3812ms	74.5616μs	13.4117 KOps/s	13.0266 KOps/s	$\color{#35bf28}+2.96\%$
test_setitem_dim[slice_int]	5.0391ms	0.1368ms	7.3099 KOps/s	7.0217 KOps/s	$\color{#35bf28}+4.10\%$
test_setitem_dim[range]	5.2982ms	0.1391ms	7.1904 KOps/s	7.3420 KOps/s	$\color{#d91a1a}-2.07\%$
test_setitem_dim[tuple]	3.1180ms	0.1136ms	8.8056 KOps/s	8.2829 KOps/s	$\textbf{\color{#35bf28}+6.31\%}$
test_setitem	66.1515μs	58.4783μs	17.1004 KOps/s	16.1878 KOps/s	$\textbf{\color{#35bf28}+5.64\%}$
test_set	78.9477μs	56.3244μs	17.7543 KOps/s	16.4452 KOps/s	$\textbf{\color{#35bf28}+7.96\%}$
test_set_shared	0.4341ms	0.3104ms	3.2221 KOps/s	3.1801 KOps/s	$\color{#35bf28}+1.32\%$
test_update	92.9990μs	69.2854μs	14.4331 KOps/s	13.4299 KOps/s	$\textbf{\color{#35bf28}+7.47\%}$
test_update_nested	0.1524ms	0.1046ms	9.5624 KOps/s	9.7522 KOps/s	$\color{#d91a1a}-1.95\%$
test_set_nested	0.1014ms	73.2682μs	13.6485 KOps/s	13.1762 KOps/s	$\color{#35bf28}+3.58\%$
test_set_nested_new	0.1452ms	96.2942μs	10.3848 KOps/s	9.0641 KOps/s	$\textbf{\color{#35bf28}+14.57\%}$
test_select	0.2089ms	0.1500ms	6.6684 KOps/s	6.3404 KOps/s	$\textbf{\color{#35bf28}+5.17\%}$
test_creation[device0]	1.3242ms	0.5884ms	1.6994 KOps/s	1.7446 KOps/s	$\color{#d91a1a}-2.59\%$
test_creation_from_tensor	0.6568ms	0.5646ms	1.7713 KOps/s	1.6506 KOps/s	$\textbf{\color{#35bf28}+7.31\%}$
test_add_one[memmap_tensor0]	0.1124ms	58.9869μs	16.9529 KOps/s	16.3667 KOps/s	$\color{#35bf28}+3.58\%$
test_contiguous[memmap_tensor0]	22.9422μs	11.4201μs	87.5649 KOps/s	84.5200 KOps/s	$\color{#35bf28}+3.60\%$
test_stack[memmap_tensor0]	0.2036ms	65.1405μs	15.3514 KOps/s	15.3468 KOps/s	$\color{#35bf28}+0.03\%$
test_reshape_pytree	0.1054ms	43.4195μs	23.0311 KOps/s	21.8370 KOps/s	$\textbf{\color{#35bf28}+5.47\%}$
test_reshape_td	0.1346ms	66.3818μs	15.0644 KOps/s	13.7984 KOps/s	$\textbf{\color{#35bf28}+9.17\%}$
test_view_pytree	50.6355μs	39.9748μs	25.0157 KOps/s	23.5857 KOps/s	$\textbf{\color{#35bf28}+6.06\%}$
test_view_td	16.2012μs	10.5608μs	94.6897 KOps/s	87.5890 KOps/s	$\textbf{\color{#35bf28}+8.11\%}$
test_unbind_pytree	0.1355ms	47.6810μs	20.9727 KOps/s	20.5839 KOps/s	$\color{#35bf28}+1.89\%$
test_unbind_td	0.3769ms	0.2894ms	3.4553 KOps/s	3.3325 KOps/s	$\color{#35bf28}+3.68\%$
test_split_pytree	64.8026μs	54.4406μs	18.3686 KOps/s	17.5695 KOps/s	$\color{#35bf28}+4.55\%$
test_split_td	0.1853ms	0.1662ms	6.0155 KOps/s	6.0617 KOps/s	$\color{#d91a1a}-0.76\%$
test_add_pytree	0.1017ms	65.5891μs	15.2464 KOps/s	15.1189 KOps/s	$\color{#35bf28}+0.84\%$
test_add_td	0.1851ms	0.1333ms	7.5005 KOps/s	7.6015 KOps/s	$\color{#d91a1a}-1.33\%$
test_distributed	68.8000μs	68.8000μs	14.5349 KOps/s	12.5313 KOps/s	$\textbf{\color{#35bf28}+15.99\%}$
test_tdmodule	3.0699ms	43.9780μs	22.7387 KOps/s	23.4828 KOps/s	$\color{#d91a1a}-3.17\%$
test_tdmodule_dispatch	60.7560ms	97.5325μs	10.2530 KOps/s	10.7976 KOps/s	$\textbf{\color{#d91a1a}-5.04\%}$
test_tdseq	2.1906ms	58.3807μs	17.1290 KOps/s	17.4849 KOps/s	$\color{#d91a1a}-2.04\%$
test_tdseq_dispatch	3.3728ms	0.1050ms	9.5197 KOps/s	9.7742 KOps/s	$\color{#d91a1a}-2.60\%$
test_instantiation_functorch	2.6146ms	2.1745ms	459.8715 Ops/s	439.7445 Ops/s	$\color{#35bf28}+4.58\%$
test_instantiation_td	8.4362ms	1.6794ms	595.4593 Ops/s	585.7320 Ops/s	$\color{#35bf28}+1.66\%$
test_exec_functorch	0.4209ms	0.3277ms	3.0513 KOps/s	3.2999 KOps/s	$\textbf{\color{#d91a1a}-7.53\%}$
test_exec_td	0.7201ms	0.5558ms	1.7993 KOps/s	1.8467 KOps/s	$\color{#d91a1a}-2.57\%$

[Blonck]

LGTM. There is just on redundant elif in _maybe_unravel_keys_silent.

[Blonck]

LGTM. There is just on redundant elif in _maybe_unravel_keys_silent.