Apply ufmt linter to all py files under test/onnx (pytorch#81335)

Same as pytorch#81285 but for `test/onnx`. The merge conflicts in `linrunner.toml` is expected. I will resolve them depending on the merge order of the PR

Pull Request resolved: pytorch#81335
Approved by: https://github.com/BowenBao, https://github.com/kit1980

Author: huydhn

.lintrunner.toml

@@ -670,6 +670,7 @@ code = 'UFMT'
 # After that, we can use **/*.py to apply the formatter to all entries
 # currently enforced by BLACK linter
 include_patterns = [
+    'test/onnx/**/*.py',
     'tools/**/*.py',
 ]
 command = [

test/onnx/debug_embed_params.py

@@ -1,9 +1,9 @@
 import sys
 
+import caffe2.python.onnx.backend as c2
+
 import onnx
 import pytorch_test_common
-
-import caffe2.python.onnx.backend as c2
 import torch
 import torch.jit
 from torch.autograd import Variable

test/onnx/export_onnx_tests_generator.py

@@ -5,10 +5,10 @@
 
 import onnx
 import onnx_test_common
-from onnx import numpy_helper
-from test_nn import new_module_tests
 
 import torch
+from onnx import numpy_helper
+from test_nn import new_module_tests
 from torch.autograd import Variable
 from torch.testing._internal.common_nn import module_tests
 

test/onnx/test_caffe2_common.py

@@ -3,12 +3,12 @@
 import glob
 import os
 
+import caffe2.python.onnx.backend as c2
+
 import numpy as np
 import onnx.backend.test
 from onnx import numpy_helper
 
-import caffe2.python.onnx.backend as c2
-
 
 def load_tensor_as_numpy_array(f):
     tensor = onnx.TensorProto()

test/onnx/test_custom_ops.py

@@ -1,13 +1,12 @@
 # Owner(s): ["module: onnx"]
 
+import caffe2.python.onnx.backend as c2
 import numpy as np
 import onnx
 import onnx_test_common
-from test_pytorch_onnx_caffe2 import do_export
-
-import caffe2.python.onnx.backend as c2
 import torch
 import torch.utils.cpp_extension
+from test_pytorch_onnx_caffe2 import do_export
 from torch.onnx import symbolic_helper
 from torch.testing._internal import common_utils
 

test/onnx/test_models.py

@@ -2,6 +2,9 @@
 
 import unittest
 
+import caffe2.python.onnx.backend as backend
+import torch
+
 from model_defs.dcgan import _netD, _netG, bsz, imgsz, nz, weights_init
 from model_defs.emb_seq import EmbeddingNetwork1, EmbeddingNetwork2
 from model_defs.mnist import MNIST
@@ -10,6 +13,11 @@
 from model_defs.srresnet import SRResNet
 from model_defs.super_resolution import SuperResolutionNet
 from pytorch_test_common import skipIfUnsupportedMinOpsetVersion, skipScriptTest
+from torch import quantization
+from torch.autograd import Variable
+from torch.onnx import OperatorExportTypes
+from torch.testing._internal import common_utils
+from torch.testing._internal.common_utils import skipIfNoLapack
 from torchvision.models import shufflenet_v2_x1_0
 from torchvision.models.alexnet import alexnet
 from torchvision.models.densenet import densenet121
@@ -23,14 +31,6 @@
 from torchvision.models.video import mc3_18, r2plus1d_18, r3d_18
 from verify import verify
 
-import caffe2.python.onnx.backend as backend
-import torch
-from torch import quantization
-from torch.autograd import Variable
-from torch.onnx import OperatorExportTypes
-from torch.testing._internal import common_utils
-from torch.testing._internal.common_utils import skipIfNoLapack
-
 if torch.cuda.is_available():
 
     def toC(x):

test/onnx/test_models_onnxruntime.py

@@ -8,9 +8,13 @@
 import onnx_test_common
 import parameterized
 import PIL
+
+import torch
 import torchvision
 from pytorch_test_common import skipIfUnsupportedMinOpsetVersion, skipScriptTest
 from test_models import TestModels
+from torch import nn
+from torch.testing._internal import common_utils
 from torchvision import ops
 from torchvision.models.detection import (
     faster_rcnn,
@@ -22,10 +26,6 @@
     transform,
 )
 
-import torch
-from torch import nn
-from torch.testing._internal import common_utils
-
 
 def exportTest(self, model, inputs, rtol=1e-2, atol=1e-7, opset_versions=None):
     opset_versions = opset_versions if opset_versions else [7, 8, 9, 10, 11, 12, 13, 14]

test/onnx/test_operators.py

@@ -8,20 +8,20 @@
 import shutil
 import tempfile
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.onnx
+
 from pytorch_test_common import (
     BATCH_SIZE,
+    flatten,
     RNN_HIDDEN_SIZE,
     RNN_INPUT_SIZE,
     RNN_SEQUENCE_LENGTH,
-    flatten,
 )
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.onnx
 from torch.autograd import Function, Variable
-from torch.nn import Module, functional
+from torch.nn import functional, Module
 from torch.onnx.symbolic_helper import (
     _get_tensor_dim_size,
     _get_tensor_sizes,

test/onnx/test_pytorch_helper.py

@@ -4,12 +4,12 @@
 import unittest
 
 import numpy as np
-from pytorch_helper import PyTorchModule
 
 import torch.nn.init as init
 import torch.onnx
 from caffe2.python.core import workspace
 from caffe2.python.model_helper import ModelHelper
+from pytorch_helper import PyTorchModule
 from torch import nn
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_utils import skipIfNoLapack

test/onnx/test_pytorch_onnx_caffe2.py

@@ -6,11 +6,20 @@
 import unittest
 from typing import Tuple
 
+import caffe2.python.onnx.backend as c2
+
 import model_defs.dcgan as dcgan
 import model_defs.word_language_model as word_language_model
 import numpy as np
 import onnx
+import torch.onnx
+import torch.onnx.operators
+import torch.utils.model_zoo as model_zoo
 import verify
+from caffe2.python.operator_test.torch_integration_test import (
+    create_bbox_transform_inputs,
+    generate_rois_rotated,
+)
 from debug_embed_params import run_embed_params
 from model_defs.lstm_flattening_result import LstmFlatteningResult
 from model_defs.mnist import MNIST
@@ -29,6 +38,12 @@
     skipIfUnsupportedMinOpsetVersion,
     skipIfUnsupportedOpsetVersion,
 )
+from torch import nn
+from torch.autograd import function, Variable
+from torch.nn.utils import rnn as rnn_utils
+from torch.onnx import ExportTypes
+from torch.testing._internal import common_utils
+from torch.testing._internal.common_utils import skipIfNoLapack
 
 # Import various models for testing
 from torchvision.models.alexnet import alexnet
@@ -37,21 +52,6 @@
 from torchvision.models.resnet import resnet50
 from torchvision.models.vgg import vgg16, vgg16_bn, vgg19, vgg19_bn
 
-import caffe2.python.onnx.backend as c2
-import torch.onnx
-import torch.onnx.operators
-import torch.utils.model_zoo as model_zoo
-from caffe2.python.operator_test.torch_integration_test import (
-    create_bbox_transform_inputs,
-    generate_rois_rotated,
-)
-from torch import nn
-from torch.autograd import Variable, function
-from torch.nn.utils import rnn as rnn_utils
-from torch.onnx import ExportTypes
-from torch.testing._internal import common_utils
-from torch.testing._internal.common_utils import skipIfNoLapack
-
 skip = unittest.skip
 
 

test/onnx/test_pytorch_onnx_caffe2_quantized.py

@@ -2,10 +2,10 @@
 
 import io
 
+import caffe2.python.onnx.backend as c2
+
 import numpy as np
 import onnx
-
-import caffe2.python.onnx.backend as c2
 import torch.nn as nn
 import torch.nn.quantized as nnq
 import torch.onnx

test/onnx/test_pytorch_onnx_no_runtime.py

@@ -6,16 +6,16 @@
 import io
 import itertools
 import unittest
-from typing import Dict, Optional, Type, Callable, Iterable, Tuple, Union, List
+from typing import Callable, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import onnx
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
-from torch.onnx import symbolic_helper, utils, symbolic_registry
+from torch.onnx import symbolic_helper, symbolic_registry, utils
 from torch.onnx._globals import GLOBALS
 from torch.testing._internal import common_utils
-import torch.nn.functional as F
 
 
 def export_to_onnx(

test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -12,6 +12,8 @@
 import numpy as np
 import onnx_test_common
 import parameterized
+
+import torch
 import torchvision
 from model_defs import (
     lstm_flattening_result,
@@ -31,8 +33,6 @@
     skipScriptTest,
     skipTraceTest,
 )
-
-import torch
 from torch import Tensor
 from torch.nn.utils import rnn as rnn_utils
 from torch.onnx import verification

test/onnx/test_pytorch_onnx_onnxruntime_cuda.py

@@ -3,15 +3,15 @@
 import unittest
 
 import onnxruntime  # noqa: F401
+
+import torch
 from pytorch_test_common import (
     skipIfNoBFloat16Cuda,
     skipIfNoCuda,
     skipIfUnsupportedMinOpsetVersion,
     skipScriptTest,
 )
 from test_pytorch_onnx_onnxruntime import TestONNXRuntime
-
-import torch
 from torch.cuda.amp import autocast
 from torch.onnx._globals import GLOBALS
 from torch.testing._internal import common_utils

test/onnx/test_pytorch_onnx_shape_inference.py

@@ -1,9 +1,9 @@
 # Owner(s): ["module: onnx"]
 
 import numpy as np
-from pytorch_test_common import skipIfUnsupportedMinOpsetVersion
 
 import torch
+from pytorch_test_common import skipIfUnsupportedMinOpsetVersion
 from torch.onnx import _constants, symbolic_helper
 from torch.testing._internal import common_utils
 

test/onnx/test_utility_funs.py

@@ -4,22 +4,21 @@
 import io
 
 import onnx
+
+import torch
+import torch.onnx
+import torch.utils.cpp_extension
 import torchvision
 from autograd_helper import CustomFunction as CustomFunction2
 from pytorch_test_common import (
     skipIfNoCuda,
     skipIfUnsupportedMaxOpsetVersion,
     skipIfUnsupportedMinOpsetVersion,
 )
-from verify import verify
-
-import torch
-import torch.onnx
-import torch.utils.cpp_extension
 from torch.onnx import (
     OperatorExportTypes,
-    TrainingMode,
     register_custom_op_symbolic,
+    TrainingMode,
     unregister_custom_op_symbolic,
     utils,
 )
@@ -31,6 +30,7 @@
     parse_args,
 )
 from torch.testing._internal import common_utils
+from verify import verify
 
 
 class _BaseTestCase(common_utils.TestCase):

test/onnx/test_verify.py

@@ -1,12 +1,11 @@
 # Owner(s): ["module: onnx"]
 
-from verify import verify
-
 import caffe2.python.onnx.backend as backend
 import torch
 from torch.autograd import Function
 from torch.nn import Module, Parameter
 from torch.testing._internal import common_utils
+from verify import verify
 
 
 class TestVerify(common_utils.TestCase):