Enable UFMT on all files in PyTorch

[ezyang]

🐛 Describe the bug

Right now, around 1.5k files are grandfathered in as "not UFMT'ed" in the PyTorch codebase. Let's make them all formatted!

To format a file, remove its name from exclude_patterns of UFMT section in .lintrunner.toml. Then, run lintrunner -a --take UFMT --all-files. You'll end up with a pull request like this: #123061

For ease of review, the UFMT diff should be 100% generatable from the .lintrunner.toml change. In particular, if you are code reviewing a UFMT change from an untrusted source, clear all changes except the .lintrunner.toml and manually rerun lintrunner to regenerate the changes. To verify that there are no changes, a maintainer can use this script: You may need to make fixup changes. This means maintainers MUST verify the additional changes on top of lint. This can be done with this script:

#!/bin/bash
set -ex
gh pr checkout $1
REV=$(git rev-parse HEAD)
git fetch origin
git checkout $(git merge-base origin/main HEAD)
git checkout $REV .lintrunner.toml
lintrunner init
lintrunner -a --take UFMT --all-files
git diff $REV

In some cases, preparatory work needs to be done. Here are some known problems:

• torch/cuda/_memory_viz.py - mypy type: ignore annotations get moved in such a way that an ignore is no longer applied to the correct line
• test/functorch/test_control_flow.py - flake8: reformatting moves flake8 ignore to the wrong line, and now a multiline string no longer has flake8 suppressed
• torch/masked/ - the import reordering causes an import cycle
• ... probably more ...

These problems should be resolved in preparatory commits BEFORE the UFMT commit, for ease of review.

How to split the work? I decided to bundle things up by directory, so long as there were not too many changes in one directory that necessitated splitting it up. With https://gist.github.com/ezyang/6c10a63d5795cbd7787beeea94497948 I did a sweep to find an optimal threshold:

13k seems like a good splitting point. Here is the worklist:

• 7 - test/_test_bazel.py
• 1831 - test/ao
• 635 - test/autograd
• 33 - test/backends
• 231 - test/benchmark_utils
• 9 - test/bottleneck_test
• 49 - test/conftest.py
• 127 - test/cpp
• 660 - test/cpp_api_parity
• 108 - test/cpp_extensions
• 4 - test/create_dummy_torchscript_model.py
• 14 - test/custom_backend
• 32 - test/custom_operator
• 6342 - test/distributed Enable UFMT on all of test/distributed #123539
• 5766 - test/distributions UFMT formatting on test/distributions, test/error_messages, test/forward_backward_compatability #123527
• 91 - test/error_messages UFMT formatting on test/distributions, test/error_messages, test/forward_backward_compatability #123527
• 837 - test/export UFMT formatting on test/export #123520 UFMT formatting on test/distributions, test/error_messages, test/forward_backward_compatability #123527
• 40 - test/forward_backward_compatibility UFMT formatting on test/distributions, test/error_messages, test/forward_backward_compatability #123527
• 11846 - test/functorch Enable UFMT on test/functorch #123541 BLOCKED
• 1784 - test/fx Enable UFMT on all files in PyTorch #123062
• 9432 - test/jit Enable UFMT on all of test/jit #123623
• 59 - test/jit_hooks
• 348 - test/lazy
• 3 - test/linear.py
• 3 - test/load_torchscript_model.py
• 5 - test/mkl_verbose.py
• 6 - test/mkldnn_verbose.py
• 665 - test/mobile UFMT formatting on test/mobile #123521
• 7397 - test/nn
• 1 - test/onnx_caffe2 Enable UFMT on test/onnx_caffe2, test/optim, test/package and test/profiler #123901
• 346 - test/optim Enable UFMT on test/onnx_caffe2, test/optim, test/package and test/profiler #123901
• 41 - test/package Enable UFMT on test/onnx_caffe2, test/optim, test/package and test/profiler #123901
• 1019 - test/profiler Enable UFMT on test/onnx_caffe2, test/optim, test/package and test/profiler #123901
• 875 - test/quantization/ao_migration Enable UFMT on all of test/quantization/ao_migration &bc #123994
• 339 - test/quantization/bc
• 12315 - test/quantization/core
• 2321 - test/quantization/eager
• 7054 - test/quantization/fx
• 544 - test/quantization/jit
• 697 - test/quantization/pt2e
• 130 - test/scripts
• 42 - test/simulate_nccl_errors.py
• 63 - test/test_ao_sparsity.py
• 134 - test/test_autocast.py
• 2917 - test/test_autograd.py
• 204 - test/test_binary_ufuncs.py
• 16 - test/test_bundled_images.py
• 114 - test/test_bundled_inputs.py
• 5 - test/test_comparison_utils.py
• 33 - test/test_compile_benchmark_util.py
• 381 - test/test_complex.py
• 55 - test/test_cpp_api_parity.py
• 57 - test/test_cpp_extensions_aot.py
• 166 - test/test_cpp_extensions_jit.py
• 190 - test/test_cpp_extensions_open_device_registration.py
• 1811 - test/test_cuda.py Enable UFMT on test/test_cuda*.py #124352
• 7 - test/test_cuda_expandable_segments.py Enable UFMT on test/test_cuda*.py #124352
• 423 - test/test_cuda_multigpu.py Enable UFMT on test/test_cuda*.py #124352
• 94 - test/test_cuda_nvml_based_avail.py Enable UFMT on test/test_cuda*.py #124352
• 41 - test/test_cuda_primary_ctx.py Enable UFMT on test/test_cuda*.py #124352
• 4 - test/test_cuda_sanitizer.py Enable UFMT on test/test_cuda*.py #124352
• 2 - test/test_cuda_trace.py Enable UFMT on test/test_cuda*.py #124352
• 1468 - test/test_dataloader.py
• 1431 - test/test_datapipe.py
• 468 - test/test_decomp.py
• 3 - test/test_deploy.py
• 17 - test/test_determination.py
• 743 - test/test_dispatch.py Enable UFMT on test/test_dispatch #123644
• 68 - test/test_dlpack.py
• 362 - test/test_dynamic_shapes.py
• 713 - test/test_expanded_weights.py
• 359 - test/test_fake_tensor.py
• 200 - test/test_flop_counter.py
• 741 - test/test_foreach.py Enable UFMT on test/test_foreach.py #123718
• 252 - test/test_function_schema.py
• 57 - test/test_functional_autograd_benchmark.py
• 24 - test/test_functional_optim.py
• 846 - test/test_functionalization.py Enable UFMT on test/test_functionalization.py #123926
• 51 - test/test_functionalization_of_rng_ops.py
• 36 - test/test_futures.py
• 1435 - test/test_fx.py
• 130 - test/test_fx_experimental.py
• 410 - test/test_fx_passes.py
• 91 - test/test_fx_reinplace_pass.py
• 179 - test/test_hub.py
• 2 - test/test_import_stats.py
• 893 - test/test_indexing.py
• 9 - test/test_itt.py
• 5290 - test/test_jit.py
• 182 - test/test_jit_autocast.py
• 20 - test/test_jit_disabled.py
• 530 - test/test_jit_fuser.py
• 3 - test/test_jit_fuser_legacy.py
• 891 - test/test_jit_fuser_te.py
• 3 - test/test_jit_legacy.py
• 343 - test/test_jit_llga_fuser.py
• 3 - test/test_jit_profiling.py
• 3 - test/test_jit_simple.py
• 229 - test/test_jit_string.py
• 87 - test/test_jiterator.py
• 34 - test/test_kernel_launch_checks.py
• 904 - test/test_legacy_vmap.py
• 28 - test/test_license.py
• 4375 - test/test_linalg.py
• 8 - test/test_logging.py
• 289 - test/test_masked.py
• 153 - test/test_maskedtensor.py
• 151 - test/test_matmul_cuda.py
• 1041 - test/test_meta.py
• 150 - test/test_metal.py
• 36 - test/test_mkl_verbose.py
• 827 - test/test_mkldnn.py
• 298 - test/test_mkldnn_fusion.py
• 36 - test/test_mkldnn_verbose.py
• 336 - test/test_mobile_optimizer.py
• 70 - test/test_model_dump.py
• 767 - test/test_modules.py
• 29 - test/test_monitor.py
• 8262 - test/test_mps.py
• 442 - test/test_multiprocessing.py Enable UFMT on test/test_multiprocessing #123840
• 59 - test/test_multiprocessing_spawn.py
• 1990 - test/test_namedtensor.py
• 393 - test/test_namedtuple_return_api.py
• 45 - test/test_native_functions.py
• 138 - test/test_native_mha.py
• 1951 - test/test_nestedtensor.py
• 9457 - test/test_nn.py
• 244 - test/test_nnapi.py
• 74 - test/test_numba_integration.py
• 179 - test/test_numpy_interop.py
• 20 - test/test_openmp.py
• 755 - test/test_ops.py Enable UFMT on test/test_ops* #123935
• 71 - test/test_ops_fwd_gradients.py
• 42 - test/test_ops_gradients.py
• 190 - test/test_ops_jit.py
• 442 - test/test_optim.py
• 115 - test/test_out_dtype_op.py
• 226 - test/test_overrides.py
• 29 - test/test_package.py
• 27 - test/test_per_overload_api.py
• 132 - test/test_prims.py
• 1016 - test/test_proxy_tensor.py
• 80 - test/test_pruning_op.py
• 159 - test/test_public_bindings.py
• 737 - test/test_python_dispatch.py
• 234 - test/test_quantization.py
• 2343 - test/test_reductions.py
• 296 - test/test_scatter_gather_ops.py
• 209 - test/test_schema_check.py
• 206 - test/test_segment_reductions.py
• 4013 - test/test_serialization.py
• 6 - test/test_set_default_mobile_cpu_allocator.py
• 239 - test/test_shape_ops.py
• 13 - test/test_show_pickle.py
• 646 - test/test_sort_and_select.py
• 3657 - test/test_sparse.py
• 3338 - test/test_sparse_csr.py
• 234 - test/test_sparse_semi_structured.py
• 1132 - test/test_spectral_ops.py
• 713 - test/test_stateless.py
• 44 - test/test_static_runtime.py
• 65 - test/test_subclass.py
• 101 - test/test_sympy_utils.py
• 2166 - test/test_tensor_creation_ops.py
• 580 - test/test_tensorboard.py
• 208 - test/test_tensorexpr.py
• 5 - test/test_tensorexpr_pybind.py
• 1189 - test/test_testing.py
• 8 - test/test_throughput_benchmark.py
• 4817 - test/test_torch.py
• 3285 - test/test_transformers.py
• 36 - test/test_type_hints.py
• 50 - test/test_type_info.py
• 606 - test/test_type_promotion.py
• 100 - test/test_unary_ufuncs.py
• 458 - test/test_utils.py
• 551 - test/test_view_ops.py
• 155 - test/test_vulkan.py
• 1365 - test/test_xnnpack_integration.py
• 175 - test/typing
• 6 - torch/_awaits
• 315 - torch/_custom_op
• 541 - torch/_export
• 2427 - torch/_functorch [BE] UFMT directory torch/_functorch #123723
• 87 - torch/_higher_order_ops
• 2 - torch/_vendor
• 2 - torch/ao/__init__.py
• 7593 - torch/ao/nn
• 2134 - torch/ao/ns
• 1442 - torch/ao/pruning
• 10772 - torch/ao/quantization
• 14 - torch/compiler
• 76 - torch/contrib
• 503 - torch/cuda
• 4071 - torch/distributed
• 233 - torch/fft
• 17 - torch/func
• 430 - torch/functional.py
• 22 - torch/futures
• 7711 - torch/fx
• 262 - torch/hub.py
• 133 - torch/library.py
• 395 - torch/linalg
• 337 - torch/masked
• 1 - torch/monitor
• 291 - torch/nested
• 10425 - torch/nn
• 3774 - torch/optim
• 476 - torch/overrides.py
• 84 - torch/quasirandom.py
• 51 - torch/random.py
• 9 - torch/return_types.py
• 729 - torch/serialization.py
• 507 - torch/signal
• 1190 - torch/sparse
• 904 - torch/special
• 675 - torch/storage.py
• 640 - torch/testing/_internal/autocast_test_lists.py
• 103 - torch/testing/_internal/autograd_function_db.py
• 32 - torch/testing/_internal/check_kernel_launches.py
• 173 - torch/testing/_internal/common_cuda.py
• 614 - torch/testing/_internal/common_device_type.py
• 124 - torch/testing/_internal/common_distributed.py
• 81 - torch/testing/_internal/common_dtype.py
• 251 - torch/testing/_internal/common_jit.py
• 9191 - torch/testing/_internal/common_modules.py
• 2889 - torch/testing/_internal/common_nn.py
• 17 - torch/testing/_internal/common_pruning.py
• 1072 - torch/testing/_internal/common_quantization.py
• 101 - torch/testing/_internal/common_quantized.py
• 68 - torch/testing/_internal/common_subclass.py
• 2289 - torch/testing/_internal/common_utils.py
• 327 - torch/testing/_internal/composite_compliance.py
• 259 - torch/testing/_internal/custom_op_db.py
• 2 - torch/testing/_internal/data
• 9 - torch/testing/_internal/dist_utils.py
• 2052 - torch/testing/_internal/distributed
• 40 - torch/testing/_internal/hop_db.py
• 160 - torch/testing/_internal/hypothesis_utils.py
• 24 - torch/testing/_internal/inductor_utils.py
• 1215 - torch/testing/_internal/jit_metaprogramming_utils.py
• 361 - torch/testing/_internal/jit_utils.py
• 88 - torch/testing/_internal/logging_tensor.py
• 32 - torch/testing/_internal/logging_utils.py
• 100 - torch/testing/_internal/optests
• 1 - torch/testing/_internal/quantization_torch_package_models.py
• 1 - torch/testing/_internal/test_module
• 14 - torch/torch_version.py Enable UFMT on torch_version.py and types.py #123131
• 27 - torch/types.py Enable UFMT on torch_version.py and types.py #123131
• 9910 - torch/utils
• 12 - torch/xpu

This list does not cover torch/testing/_internal/common_methods_invocations.py, which has a whopping 35341 lines of changes and will likely have to be handled in some other way (e.g., splitting it first).

To claim an entry, put your name (and eventually PR) next to it by editing this post. Feel free to bundle multiple small but related changes together if necessary.

Versions

main

[Skylion007]

@ezyang While we are doing this, we should investigate using ruff isort rules + ruff fmt to remove the need for ufmt in the future. Definitely, I would focus on getting full coverage by ruff format first though.

Benefits:

• increase the speed of the formatting changes / make them fully integratable with vscode
• reduces dependencies
• more configurable than black

Downsides

• ruff is not as stable as black
• ties fmtting version to linter versioning
• complicates ignore logic for ruff potentially (we need to selectively enable isort in places we may not care about other linting)

[ezyang]

Way faster than black, sign me up!

@Skylion007 if we are actually serious about this, we should flag day from UFMT to Ruff first (for all the files that are already formatted), and pause UFMT'ing until we can single step to ruff fmt instead. Are you volunteering to do this? lol

[XuehaiPan]

@ezyang I think we should make all linters and formatters have the same line-length first.

• Add pyproject.toml for black configuration #79399 (comment)

We are now using the default line length from black which is 88.

pytorch/pyproject.toml

Lines 17 to 25 in 270dd99

	[tool.black]
	# Uncomment if pyproject.toml worked fine to ensure consistency with flake8
	# line-length = 120
	target-version = ["py38", "py39", "py310", "py311"]
	[tool.ruff]
	target-version = "py38"
	line-length = 120

pytorch/.flake8

Lines 1 to 6 in 270dd99

	[flake8]
	# NOTE: **Mirror any changes** to this file the [tool.ruff] config in pyproject.toml
	# before we can fully move to use ruff
	enable-extensions = G
	select = B,C,E,F,G,P,SIM1,T4,W,B9,TOR0,TOR1,TOR2,TOR9
	max-line-length = 120

Personally, I think line-length = 100 would be a reasonable value. But any change to the black line length configuration would create a huge PR. This would create a huge PR for unwrapping the multi-line expression into a single line if it fits or vice versa.

[XuehaiPan]

A longer line-length might be needed for #121636.

• [BE]: Update ruff to 0.3.2 #121636

Because ruff tries to convert str.format call to f-string, which might be very long.

- "{} some text {}".format(
-     very_very_very_very_very_very_very_very_very_long_expression,
-     very_long_long_long_long_long_long_long_long_long_expression,
- )
+ f"{very_very_very_very_very_very_very_very_very_long_expression} some text {very_long_long_long_long_long_long_long_long_long_expression}"

[Chubercik]

Personally, I think line-length = 100 would be a reasonable value. But any change to the black line length configuration would create a huge PR. This would create a huge PR for unwrapping the multi-line expression into a single line if it fits or vice versa.

Agreed, and I think it'd be worth it in the end; if we decide to go forward with that, a similar tracking issue to this one could be posted up, and it could be done in the same manner we are UFMT-ing right now.

[ezyang]

@XuehaiPan oof, what a pain. I think 88 is what we're using inside Meta internally, so I think my preference is to stick with it. It seems like the flake8 thing is to stop flake8 from bleating when black is unable to get it quite short enough.

[malfet]

If I remember discussion correctly from a few years back, tests are not part of PyTorch package and as such should not necessarily abide by the same rules as rest of the codebase. I.e. linting tests are fine, but imo optional.

[ezyang]

@malfet I'd like tests to be ufmt'ed, primarily because it's nice to botch my formatting with lines too long and then have ufmt fix it. If ufmt is optional, then arguably flake8 line length should be optional too :P