[Lint] Add TorchFix linter

[kit1980]

Add TorchFix (https://github.com/pytorch/test-infra/tree/main/tools/torchfix) flake8 plugin.

To make it pass:

1. Exclude files where TorchFix finds functorch violations as older version still needs to be supported
2. Update deprecated torch.norm to torch.linalg.norm.

[github-actions]

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

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

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_single	0.1290s	0.1187s	8.4251 Ops/s	8.2903 Ops/s	$\color{#35bf28}+1.63\%$
test_sync	78.8146ms	66.2323ms	15.0984 Ops/s	15.5310 Ops/s	$\color{#d91a1a}-2.79\%$
test_async	0.1275s	63.2427ms	15.8121 Ops/s	12.8271 Ops/s	$\textbf{\color{#35bf28}+23.27\%}$
test_simple	1.1123s	1.0167s	0.9836 Ops/s	0.9394 Ops/s	$\color{#35bf28}+4.70\%$
test_transformed	1.3750s	1.2753s	0.7842 Ops/s	0.7541 Ops/s	$\color{#35bf28}+3.99\%$
test_serial	2.9826s	2.9471s	0.3393 Ops/s	0.3346 Ops/s	$\color{#35bf28}+1.40\%$
test_parallel	2.6302s	2.5370s	0.3942 Ops/s	0.3946 Ops/s	$\color{#d91a1a}-0.10\%$
test_step_mdp_speed[True-True-True-True-True]	2.3582ms	62.1375μs	16.0933 KOps/s	15.1320 KOps/s	$\textbf{\color{#35bf28}+6.35\%}$
test_step_mdp_speed[True-True-True-True-False]	3.5034ms	37.2575μs	26.8403 KOps/s	27.1897 KOps/s	$\color{#d91a1a}-1.29\%$
test_step_mdp_speed[True-True-True-False-True]	2.5997ms	45.8412μs	21.8145 KOps/s	21.8438 KOps/s	$\color{#d91a1a}-0.13\%$
test_step_mdp_speed[True-True-True-False-False]	4.1666ms	25.1952μs	39.6902 KOps/s	38.7311 KOps/s	$\color{#35bf28}+2.48\%$
test_step_mdp_speed[True-True-False-True-True]	3.0044ms	65.6775μs	15.2259 KOps/s	15.2063 KOps/s	$\color{#35bf28}+0.13\%$
test_step_mdp_speed[True-True-False-True-False]	1.9967ms	40.4261μs	24.7365 KOps/s	24.3984 KOps/s	$\color{#35bf28}+1.39\%$
test_step_mdp_speed[True-True-False-False-True]	5.4125ms	48.5048μs	20.6165 KOps/s	19.5568 KOps/s	$\textbf{\color{#35bf28}+5.42\%}$
test_step_mdp_speed[True-True-False-False-False]	3.6930ms	26.7018μs	37.4507 KOps/s	35.2792 KOps/s	$\textbf{\color{#35bf28}+6.16\%}$
test_step_mdp_speed[True-False-True-True-True]	3.9456ms	69.5467μs	14.3788 KOps/s	14.3367 KOps/s	$\color{#35bf28}+0.29\%$
test_step_mdp_speed[True-False-True-True-False]	3.7245ms	40.1722μs	24.8928 KOps/s	24.4270 KOps/s	$\color{#35bf28}+1.91\%$
test_step_mdp_speed[True-False-True-False-True]	9.6516ms	49.2670μs	20.2976 KOps/s	19.4030 KOps/s	$\color{#35bf28}+4.61\%$
test_step_mdp_speed[True-False-True-False-False]	2.0545ms	27.3085μs	36.6187 KOps/s	35.2751 KOps/s	$\color{#35bf28}+3.81\%$
test_step_mdp_speed[True-False-False-True-True]	4.6801ms	70.9644μs	14.0916 KOps/s	13.5849 KOps/s	$\color{#35bf28}+3.73\%$
test_step_mdp_speed[True-False-False-True-False]	5.6910ms	42.6206μs	23.4628 KOps/s	23.5825 KOps/s	$\color{#d91a1a}-0.51\%$
test_step_mdp_speed[True-False-False-False-True]	3.8482ms	49.0295μs	20.3959 KOps/s	20.2679 KOps/s	$\color{#35bf28}+0.63\%$
test_step_mdp_speed[True-False-False-False-False]	1.9758ms	29.8862μs	33.4603 KOps/s	33.3899 KOps/s	$\color{#35bf28}+0.21\%$
test_step_mdp_speed[False-True-True-True-True]	3.7946ms	67.8714μs	14.7338 KOps/s	14.8857 KOps/s	$\color{#d91a1a}-1.02\%$
test_step_mdp_speed[False-True-True-True-False]	7.3432ms	42.9740μs	23.2699 KOps/s	24.4138 KOps/s	$\color{#d91a1a}-4.69\%$
test_step_mdp_speed[False-True-True-False-True]	1.7514ms	50.4713μs	19.8133 KOps/s	18.6920 KOps/s	$\textbf{\color{#35bf28}+6.00\%}$
test_step_mdp_speed[False-True-True-False-False]	3.0473ms	30.3389μs	32.9609 KOps/s	31.3876 KOps/s	$\textbf{\color{#35bf28}+5.01\%}$
test_step_mdp_speed[False-True-False-True-True]	2.0088ms	67.5241μs	14.8095 KOps/s	14.3999 KOps/s	$\color{#35bf28}+2.84\%$
test_step_mdp_speed[False-True-False-True-False]	3.7054ms	42.5749μs	23.4880 KOps/s	23.3056 KOps/s	$\color{#35bf28}+0.78\%$
test_step_mdp_speed[False-True-False-False-True]	3.1162ms	53.2552μs	18.7775 KOps/s	18.3784 KOps/s	$\color{#35bf28}+2.17\%$
test_step_mdp_speed[False-True-False-False-False]	2.9379ms	32.3980μs	30.8662 KOps/s	30.2509 KOps/s	$\color{#35bf28}+2.03\%$
test_step_mdp_speed[False-False-True-True-True]	3.3607ms	70.3043μs	14.2239 KOps/s	13.6296 KOps/s	$\color{#35bf28}+4.36\%$
test_step_mdp_speed[False-False-True-True-False]	2.3197ms	44.8569μs	22.2931 KOps/s	22.0153 KOps/s	$\color{#35bf28}+1.26\%$
test_step_mdp_speed[False-False-True-False-True]	1.9451ms	53.9594μs	18.5324 KOps/s	19.1108 KOps/s	$\color{#d91a1a}-3.03\%$
test_step_mdp_speed[False-False-True-False-False]	2.9390ms	33.1348μs	30.1797 KOps/s	31.1260 KOps/s	$\color{#d91a1a}-3.04\%$
test_step_mdp_speed[False-False-False-True-True]	1.9771ms	72.7611μs	13.7436 KOps/s	13.3920 KOps/s	$\color{#35bf28}+2.63\%$
test_step_mdp_speed[False-False-False-True-False]	2.0056ms	46.2748μs	21.6100 KOps/s	20.8893 KOps/s	$\color{#35bf28}+3.45\%$
test_step_mdp_speed[False-False-False-False-True]	3.7509ms	57.8271μs	17.2929 KOps/s	17.9240 KOps/s	$\color{#d91a1a}-3.52\%$
test_step_mdp_speed[False-False-False-False-False]	4.8407ms	34.6090μs	28.8943 KOps/s	28.9520 KOps/s	$\color{#d91a1a}-0.20\%$
test_values[generalized_advantage_estimate-True-True]	33.9481ms	22.7701ms	43.9172 Ops/s	45.7193 Ops/s	$\color{#d91a1a}-3.94\%$
test_values[vec_generalized_advantage_estimate-True-True]	92.6579ms	82.0728ms	12.1843 Ops/s	12.3163 Ops/s	$\color{#d91a1a}-1.07\%$
test_values[td0_return_estimate-False-False]	2.6288ms	0.9569ms	1.0450 KOps/s	1.1482 KOps/s	$\textbf{\color{#d91a1a}-8.99\%}$
test_values[td1_return_estimate-False-False]	26.0344ms	21.2208ms	47.1236 Ops/s	47.9136 Ops/s	$\color{#d91a1a}-1.65\%$
test_values[vec_td1_return_estimate-False-False]	99.5963ms	82.2561ms	12.1572 Ops/s	12.1898 Ops/s	$\color{#d91a1a}-0.27\%$
test_values[td_lambda_return_estimate-True-False]	63.9084ms	55.6586ms	17.9667 Ops/s	17.6272 Ops/s	$\color{#35bf28}+1.93\%$
test_values[vec_td_lambda_return_estimate-True-False]	89.4125ms	81.7560ms	12.2315 Ops/s	12.2128 Ops/s	$\color{#35bf28}+0.15\%$
test_gae_speed[generalized_advantage_estimate-False-1-512]	20.5041ms	16.7272ms	59.7828 Ops/s	58.8645 Ops/s	$\color{#35bf28}+1.56\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512]	14.6566ms	6.6074ms	151.3464 Ops/s	163.1304 Ops/s	$\textbf{\color{#d91a1a}-7.22\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512]	9.5707ms	1.0066ms	993.4312 Ops/s	998.7191 Ops/s	$\color{#d91a1a}-0.53\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512]	0.1007s	79.9168ms	12.5130 Ops/s	12.1270 Ops/s	$\color{#35bf28}+3.18\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512]	14.8329ms	7.8151ms	127.9572 Ops/s	132.1228 Ops/s	$\color{#d91a1a}-3.15\%$
test_dqn_speed	13.4287ms	4.6109ms	216.8771 Ops/s	212.6333 Ops/s	$\color{#35bf28}+2.00\%$
test_ddpg_speed	16.1777ms	8.4531ms	118.2996 Ops/s	108.2480 Ops/s	$\textbf{\color{#35bf28}+9.29\%}$
test_sac_speed	34.8908ms	24.8337ms	40.2678 Ops/s	40.9865 Ops/s	$\color{#d91a1a}-1.75\%$
test_redq_speed	42.6702ms	34.2481ms	29.1987 Ops/s	26.8626 Ops/s	$\textbf{\color{#35bf28}+8.70\%}$
test_redq_deprec_speed	48.4303ms	34.0897ms	29.3344 Ops/s	27.2794 Ops/s	$\textbf{\color{#35bf28}+7.53\%}$
test_td3_speed	31.5485ms	24.5926ms	40.6627 Ops/s	40.7502 Ops/s	$\color{#d91a1a}-0.21\%$
test_cql_speed	97.7285ms	81.1307ms	12.3258 Ops/s	12.8199 Ops/s	$\color{#d91a1a}-3.85\%$
test_a2c_speed	44.8203ms	20.1598ms	49.6036 Ops/s	55.5660 Ops/s	$\textbf{\color{#d91a1a}-10.73\%}$
test_ppo_speed	36.1037ms	20.8741ms	47.9063 Ops/s	54.6589 Ops/s	$\textbf{\color{#d91a1a}-12.35\%}$
test_reinforce_speed	29.6406ms	16.2648ms	61.4826 Ops/s	66.9595 Ops/s	$\textbf{\color{#d91a1a}-8.18\%}$
test_iql_speed	81.1686ms	68.2420ms	14.6537 Ops/s	14.1303 Ops/s	$\color{#35bf28}+3.70\%$
test_sample_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	8.9034ms	4.1408ms	241.4985 Ops/s	256.4189 Ops/s	$\textbf{\color{#d91a1a}-5.82\%}$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	0.2662s	5.0475ms	198.1187 Ops/s	247.0520 Ops/s	$\textbf{\color{#d91a1a}-19.81\%}$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	0.2041s	4.8203ms	207.4559 Ops/s	242.9061 Ops/s	$\textbf{\color{#d91a1a}-14.59\%}$
test_sample_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	8.7525ms	4.0718ms	245.5940 Ops/s	239.8969 Ops/s	$\color{#35bf28}+2.37\%$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	10.3062ms	4.2127ms	237.3777 Ops/s	246.3911 Ops/s	$\color{#d91a1a}-3.66\%$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	6.2603ms	4.1014ms	243.8176 Ops/s	241.3109 Ops/s	$\color{#35bf28}+1.04\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	9.9820ms	4.0287ms	248.2213 Ops/s	251.6316 Ops/s	$\color{#d91a1a}-1.36\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	8.4976ms	4.0308ms	248.0881 Ops/s	238.9837 Ops/s	$\color{#35bf28}+3.81\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	8.7067ms	4.0584ms	246.4049 Ops/s	240.1654 Ops/s	$\color{#35bf28}+2.60\%$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	8.5383ms	4.1595ms	240.4141 Ops/s	249.4409 Ops/s	$\color{#d91a1a}-3.62\%$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	9.9306ms	4.3500ms	229.8864 Ops/s	229.1449 Ops/s	$\color{#35bf28}+0.32\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	10.3111ms	4.2899ms	233.1050 Ops/s	243.5365 Ops/s	$\color{#d91a1a}-4.28\%$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	6.2577ms	3.9138ms	255.5030 Ops/s	257.2228 Ops/s	$\color{#d91a1a}-0.67\%$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	8.3853ms	4.0780ms	245.2156 Ops/s	240.6498 Ops/s	$\color{#35bf28}+1.90\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	0.2051s	4.8951ms	204.2868 Ops/s	233.5190 Ops/s	$\textbf{\color{#d91a1a}-12.52\%}$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	9.8772ms	3.9438ms	253.5635 Ops/s	246.5981 Ops/s	$\color{#35bf28}+2.82\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	8.3959ms	4.1563ms	240.5977 Ops/s	243.6208 Ops/s	$\color{#d91a1a}-1.24\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	12.5196ms	4.2568ms	234.9180 Ops/s	239.2949 Ops/s	$\color{#d91a1a}-1.83\%$
test_populate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-400]	0.3931s	48.6271ms	20.5646 Ops/s	21.7617 Ops/s	$\textbf{\color{#d91a1a}-5.50\%}$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]	0.2123s	40.8053ms	24.5066 Ops/s	22.7367 Ops/s	$\textbf{\color{#35bf28}+7.78\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]	0.2076s	43.0987ms	23.2026 Ops/s	24.2219 Ops/s	$\color{#d91a1a}-4.21\%$
test_populate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]	0.2127s	39.7435ms	25.1613 Ops/s	21.7858 Ops/s	$\textbf{\color{#35bf28}+15.49\%}$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]	0.2347s	45.3617ms	22.0450 Ops/s	23.3088 Ops/s	$\textbf{\color{#d91a1a}-5.42\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]	0.2248s	41.2380ms	24.2495 Ops/s	20.9185 Ops/s	$\textbf{\color{#35bf28}+15.92\%}$
test_populate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]	0.2147s	43.7101ms	22.8780 Ops/s	23.8351 Ops/s	$\color{#d91a1a}-4.02\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]	0.2293s	41.9101ms	23.8606 Ops/s	22.0365 Ops/s	$\textbf{\color{#35bf28}+8.28\%}$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]	0.2240s	45.5232ms	21.9668 Ops/s	23.5824 Ops/s	$\textbf{\color{#d91a1a}-6.85\%}$

[kit1980]

"Continuous Benchmark" fails unrelated to this PR:

ImportError: cannot import name 'unravel_key' from 'tensordict'

[vmoens]

LGTM thanks!

[vmoens]

The olddeps tests are failing because vmap cannot be found.
We try to be compatible with older versions of pytorch to support users that rely on these (for instance, libraries that permit deployment of pytorch model on hardware and that depend on a specific pytorch version).
Any idea how we can solve that issue?

[kit1980]

@vmoens Thanks for looking.
I think I'll just exclude the files with import functorch from the linter and revert the changes to them.

[kit1980]

@vmoens I've updated the PR to ignore functorch.
Please take another look.
Eventually I want to add a configuration to specify PyTorch version for TorchFix, but it's not there yet.

[vmoens]

LGTM

[kit1980]

Hi @vmoens
I wonder if you still want to support import functorch?
Functorch moved into torch.func long time ago.
RL readme says "install the latest (nightly) PyTorch release or the latest stable version of PyTorch", and both of them of course have torch.func.

[vmoens]

Unfortunately functorch.dim is still a thing that we need and it hasn't moved to core.
But a bunch of people from core (including myself) are working on productionalizing it and porting it where it belongs!

That being said it's something users should not care about since functorch is packed with PyTorch...