ai-edge-torch-nightly 0.2.0.dev20240730__py3-none-any.whl → 0.2.0.dev20240802__py3-none-any.whl

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -16,13 +16,6 @@ import operator
 import os
 from typing import Optional, Tuple, Union
 
-import torch
-import torch.ao.quantization.quantize_pt2e
-from torch.export import ExportedProgram
-from torch.fx import GraphModule
-from torch.fx import Node
-import torch.utils._pytree as pytree
-
 from ai_edge_torch.convert.fx_passes import ExportedProgramPassBase
 from ai_edge_torch.convert.fx_passes import ExportedProgramPassResult
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
@@ -30,6 +23,12 @@ from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layo
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
+import torch
+import torch.ao.quantization.quantize_pt2e
+from torch.export import ExportedProgram
+from torch.fx import GraphModule
+from torch.fx import Node
+import torch.utils._pytree as pytree
 
 TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
 
@@ -208,7 +207,8 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
 
     if not layout_check.is_4d(input_node):
       raise AssertionError(
-          f"Attempting to convert non-NHWC compatible node to NHWC: {input_node}"
+          "Attempting to convert non-NHWC compatible node to NHWC:"
+          f" {input_node}"
       )
 
     # Assign target node's source meta to the to_NHWC node, because the transpose
@@ -250,7 +250,9 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
 
     for node in graph.nodes:
       has_input_nodes = len(node.all_input_nodes) > 0
-      all_inputs_are_const = all(map(layout_mark.is_const_node, node.all_input_nodes))
+      all_inputs_are_const = all(
+          map(layout_mark.is_const_node, node.all_input_nodes)
+      )
       if (
           node.name in non_user_input_names
           or (has_input_nodes and all_inputs_are_const)
@@ -262,7 +264,9 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    partitioner = os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None)
+    partitioner = os.environ.get(
+        "AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None
+    )
     if partitioner == "MINCUT":
       graph_module = layout_partitioners.min_cut.partition(graph_module)
     elif partitioner == "GREEDY":

ai_edge_torch/convert/test/test_convert.py

@@ -20,15 +20,14 @@ import tempfile
 from typing import Tuple
 import unittest
 
+import ai_edge_torch
+from ai_edge_torch.convert import conversion_utils as cutils
+from ai_edge_torch.testing import model_coverage
 import numpy as np
 import tensorflow as tf
 import torch
 import torchvision
 
-import ai_edge_torch
-from ai_edge_torch.convert import conversion_utils as cutils
-from ai_edge_torch.testing import model_coverage
-
 
 @dataclass
 class TestContainer1:
@@ -36,7 +35,9 @@ class TestContainer1:
   data_2: Tuple[torch.Tensor, torch.Tensor]
 
 
-torch.export.register_dataclass(TestContainer1, serialized_type_name="TestContainer1")
+torch.export.register_dataclass(
+    TestContainer1, serialized_type_name="TestContainer1"
+)
 
 
 class TestConvert(unittest.TestCase):
@@ -60,7 +61,9 @@ class TestConvert(unittest.TestCase):
     torch_module = Add().eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
   def test_convert_dot_add(self):
     class DotAdd(torch.nn.Module):
@@ -77,14 +80,18 @@ class TestConvert(unittest.TestCase):
     torch_module = DotAdd().eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
   def test_convert_resnet18(self):
     args = (torch.randn(4, 3, 224, 224),)
     torch_module = torchvision.models.resnet18().eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
   def test_signature_args_ordering(self):
     """Tests conversion of a model with more than 10 arguments."""
@@ -156,7 +163,9 @@ class TestConvert(unittest.TestCase):
     torch_model = BasicAddModelWithMultipleOutputs().eval()
     edge_model = ai_edge_torch.convert(torch_model, sample_input)
 
-    result = model_coverage.compare_tflite_torch(edge_model, torch_model, sample_input)
+    result = model_coverage.compare_tflite_torch(
+        edge_model, torch_model, sample_input
+    )
     self.assertTrue(result)
 
   def test_12_outputs_model(self):
@@ -201,7 +210,9 @@ class TestConvert(unittest.TestCase):
     torch_model = BasicAddModelWithMultipleOutputs().eval()
     edge_model = ai_edge_torch.convert(torch_model, sample_input)
 
-    result = model_coverage.compare_tflite_torch(edge_model, torch_model, sample_input)
+    result = model_coverage.compare_tflite_torch(
+        edge_model, torch_model, sample_input
+    )
     self.assertTrue(result)
 
   def test_apply_tfl_backdoor_flags(self):
@@ -244,7 +255,9 @@ class TestConvert(unittest.TestCase):
           tmp_dir_path, "test_convert_add_backdoor_flags_mlir_dump"
       )
       ai_edge_torch.convert(
-          torch_module, args, _ai_edge_converter_flags={"ir_dump_dir": ir_dump_path}
+          torch_module,
+          args,
+          _ai_edge_converter_flags={"ir_dump_dir": ir_dump_path},
       )
       self.assertTrue(os.path.isdir(ir_dump_path))
 
@@ -296,7 +309,9 @@ class TestConvert(unittest.TestCase):
     edge_model = ai_edge_torch.convert(model, sample_kwargs=kwargs_gen())
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, model, kwargs=kwargs_gen)
+        model_coverage.compare_tflite_torch(
+            edge_model, model, kwargs=kwargs_gen
+        )
     )
 
   def test_convert_model_with_args_kwargs(self):
@@ -316,7 +331,9 @@ class TestConvert(unittest.TestCase):
     edge_model = ai_edge_torch.convert(model, args_gen(), kwargs_gen())
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, model, args_gen, kwargs_gen)
+        model_coverage.compare_tflite_torch(
+            edge_model, model, args_gen, kwargs_gen
+        )
     )
 
   def test_convert_model_with_args_nested_kwargs_1(self):
@@ -344,7 +361,9 @@ class TestConvert(unittest.TestCase):
         "z_data_2_0": kwargs["z"].data_2[0].numpy(),
         "z_data_2_1": kwargs["z"].data_2[1].numpy(),
     }
-    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+    self._compare_tflite_torch_args_kwargs(
+        SampleModel(), args, kwargs, flat_inputs
+    )
 
   def test_convert_model_with_args_nested_kwargs_2(self):
     """
@@ -371,7 +390,9 @@ class TestConvert(unittest.TestCase):
         "z_data_2_0_0": kwargs["z"].data_2[0][0].numpy(),
         "z_data_2_1": kwargs["z"].data_2[1].numpy(),
     }
-    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+    self._compare_tflite_torch_args_kwargs(
+        SampleModel(), args, kwargs, flat_inputs
+    )
 
   def test_convert_model_with_args_nested_kwargs_3(self):
     """
@@ -398,7 +419,9 @@ class TestConvert(unittest.TestCase):
         "z_data_2_0_foo": kwargs["z"].data_2[0]["foo"].numpy(),
         "z_data_2_1": kwargs["z"].data_2[1].numpy(),
     }
-    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+    self._compare_tflite_torch_args_kwargs(
+        SampleModel(), args, kwargs, flat_inputs
+    )
 
   def _compare_tflite_torch_args_kwargs(self, model, args, kwargs, flat_inputs):
     model.eval()

ai_edge_torch/convert/test/test_convert_composites.py

@@ -17,11 +17,10 @@
 from typing import Callable
 import unittest
 
-import parameterized
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
+import parameterized
+import torch
 
 
 def _func_to_torch_module(func: Callable):
@@ -47,35 +46,35 @@ class TestConvertComposites(unittest.TestCase):
     torch_module = torch.nn.Hardswish().eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
-
-  @parameterized.parameterized.expand(
-      [
-          # input_size, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override
-          # no padding, stride = 1
-          ([1, 3, 6, 6], [3, 3], [1, 1], [0, 0], False, True, None),
-          # add stride
-          ([1, 3, 6, 6], [3, 3], [2, 2], [0, 0], False, True, None),
-          # default values
-          ([1, 3, 6, 6], [3, 3]),
-          # add padding
-          ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, True, None),
-          # add different padding for different dims
-          ([1, 3, 6, 6], [3, 3], [1, 1], [0, 1], False, True, None),
-          # add both stride and padding
-          ([1, 3, 6, 6], [3, 3], [2, 2], [1, 1], False, True, None),
-          # count_include_pad = False
-          ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, False, None),
-          # ceil_mode = True
-          ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], True, True, None),
-          # ceil_mode = True, stride=[3, 3]
-          ([1, 3, 6, 6], [3, 3], [3, 3], [1, 1], True, True, None),
-          # set divisor_override
-          ([1, 3, 6, 6], [3, 3], [1, 1], 0, False, True, 6),
-          # padding set to one number
-          ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
-      ]
-  )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
+  @parameterized.parameterized.expand([
+      # input_size, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override
+      # no padding, stride = 1
+      ([1, 3, 6, 6], [3, 3], [1, 1], [0, 0], False, True, None),
+      # add stride
+      ([1, 3, 6, 6], [3, 3], [2, 2], [0, 0], False, True, None),
+      # default values
+      ([1, 3, 6, 6], [3, 3]),
+      # add padding
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, True, None),
+      # add different padding for different dims
+      ([1, 3, 6, 6], [3, 3], [1, 1], [0, 1], False, True, None),
+      # add both stride and padding
+      ([1, 3, 6, 6], [3, 3], [2, 2], [1, 1], False, True, None),
+      # count_include_pad = False
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, False, None),
+      # ceil_mode = True
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], True, True, None),
+      # ceil_mode = True, stride=[3, 3]
+      ([1, 3, 6, 6], [3, 3], [3, 3], [1, 1], True, True, None),
+      # set divisor_override
+      ([1, 3, 6, 6], [3, 3], [1, 1], 0, False, True, 6),
+      # padding set to one number
+      ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
+  ])
   def test_convert_avg_pool2d(self, input_size, *args):
     """Tests conversion of a module containing an avg_pool2d aten."""
     torch_module = _func_to_torch_module(
@@ -85,51 +84,65 @@ class TestConvertComposites(unittest.TestCase):
     edge_model = ai_edge_torch.convert(torch_module, tracing_args)
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
     )
 
-  @parameterized.parameterized.expand(
-      [
-          # use scale_factor with align_corners=False
-          (
-              [1, 3, 10, 10],
-              dict(scale_factor=3.0, mode='bilinear', align_corners=False),
-          ),
-          # use scale_factor with align_corners=true
-          ([1, 3, 10, 10], dict(scale_factor=3.0, mode='bilinear', align_corners=True)),
-          # use size
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
-          # use size with align_corners=true
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear', align_corners=True)),
-      ]
-  )
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
   def test_convert_upsample_bilinear_functional(self, input_size, kwargs):
     """Tests conversion of a torch.nn.functional.upsample module."""
     torch_module = _func_to_torch_module(
-        lambda input_tensor: torch.nn.functional.upsample(input_tensor, **kwargs)
+        lambda input_tensor: torch.nn.functional.upsample(
+            input_tensor, **kwargs
+        )
     )
     tracing_args = (torch.randn(*input_size),)
     edge_model = ai_edge_torch.convert(torch_module, tracing_args)
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
     )
 
-  @parameterized.parameterized.expand(
-      [
-          # use scale_factor with align_corners=False
-          (
-              [1, 3, 10, 10],
-              dict(scale_factor=3.0, mode='bilinear', align_corners=False),
-          ),
-          # use scale_factor with align_corners=true
-          ([1, 3, 10, 10], dict(scale_factor=3.0, mode='bilinear', align_corners=True)),
-          # use size
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
-          # use size with align_corners=true
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear', align_corners=True)),
-      ]
-  )
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
   def test_convert_upsample_bilinear(self, input_size, kwargs):
     """Tests conversion of a torch.nn.Upsample module."""
     torch_module = _func_to_torch_module(
@@ -139,34 +152,44 @@ class TestConvertComposites(unittest.TestCase):
     edge_model = ai_edge_torch.convert(torch_module, tracing_args)
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
     )
 
-  @parameterized.parameterized.expand(
-      [
-          # use scale_factor with align_corners=False
-          (
-              [1, 3, 10, 10],
-              dict(scale_factor=3.0, mode='bilinear', align_corners=False),
-          ),
-          # use scale_factor with align_corners=true
-          ([1, 3, 10, 10], dict(scale_factor=3.0, mode='bilinear', align_corners=True)),
-          # use size
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
-          # use size with align_corners=true
-          ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear', align_corners=True)),
-      ]
-  )
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
   def test_convert_interpolate_bilinear_functional(self, input_size, kwargs):
     """Tests conversion of a torch.nn.functional.interpolate module."""
     torch_module = _func_to_torch_module(
-        lambda input_tensor: torch.nn.functional.interpolate(input_tensor, **kwargs)
+        lambda input_tensor: torch.nn.functional.interpolate(
+            input_tensor, **kwargs
+        )
     )
     tracing_args = (torch.randn(*input_size),)
     edge_model = ai_edge_torch.convert(torch_module, tracing_args)
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
     )
 
   def test_convert_gelu(self):
@@ -176,7 +199,9 @@ class TestConvertComposites(unittest.TestCase):
     torch_module = torch.nn.GELU().eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
   def test_convert_gelu_approximate(self):
     """Tests conversion of an Approximate GELU module."""
@@ -185,7 +210,9 @@ class TestConvertComposites(unittest.TestCase):
     torch_module = torch.nn.GELU('tanh').eval()
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
   def test_convert_embedding_lookup(self):
     """Tests conversion of an Embedding module."""
@@ -194,7 +221,9 @@ class TestConvertComposites(unittest.TestCase):
     torch_module = torch.nn.Embedding(10, 10)
     edge_model = ai_edge_torch.convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
 
 
 if __name__ == '__main__':

ai_edge_torch/convert/test/test_convert_multisig.py

@@ -15,11 +15,10 @@
 
 import unittest
 
-import torch
-import torchvision
-
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
+import torch
+import torchvision
 
 
 class TestConvertMultiSignature(unittest.TestCase):
@@ -41,7 +40,9 @@ class TestConvertMultiSignature(unittest.TestCase):
         signature_name, torch_module, large_args
     ).convert(torch_module, args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model, torch_module, large_args, signature_name=signature_name
@@ -74,7 +75,10 @@ class TestConvertMultiSignature(unittest.TestCase):
     )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
-            edge_model, torch_module, large_args, signature_name=signature_name_2
+            edge_model,
+            torch_module,
+            large_args,
+            signature_name=signature_name_2,
         )
     )
 
@@ -87,11 +91,13 @@ class TestConvertMultiSignature(unittest.TestCase):
 
     signature_name = "large_input"
 
-    edge_model = ai_edge_torch.signature(signature_name, torch_module, args).convert(
-        torch_module, large_args
-    )
+    edge_model = ai_edge_torch.signature(
+        signature_name, torch_module, args
+    ).convert(torch_module, large_args)
 
-    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model, torch_module, large_args, signature_name=signature_name
@@ -110,7 +116,9 @@ class TestConvertMultiSignature(unittest.TestCase):
     resnet_signature_name = "resnet18"
 
     edge_model = (
-        ai_edge_torch.signature(mobilenet_signature_name, mobilentv2, mobilenet_args)
+        ai_edge_torch.signature(
+            mobilenet_signature_name, mobilentv2, mobilenet_args
+        )
         .signature(resnet_signature_name, resnet18, resnet_args)
         .convert(resnet18, resnet_args)
     )

ai_edge_torch/convert/test/test_to_channel_last_io.py

@@ -15,9 +15,8 @@
 
 import unittest
 
-import torch
-
 import ai_edge_torch
+import torch
 
 
 class Identity(torch.nn.Module):

ai_edge_torch/convert/to_channel_last_io.py

@@ -31,7 +31,9 @@ class ChannelLastIOWrapper(nn.Module):
     if not torch.is_tensor(x):
       raise ValueError("Input must be a torch tensor")
     if x.ndim < 3:
-      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, C, ...)")
+      raise ValueError(
+          "Input must be a tensor with rank >= 3 in layout (N, C, ...)"
+      )
     dims = [0, *range(2, x.ndim), 1]
     return torch.permute(x, dims)
 
@@ -39,7 +41,9 @@ class ChannelLastIOWrapper(nn.Module):
     if not torch.is_tensor(x):
       raise ValueError("Input must be a torch tensor.")
     if x.ndim < 3:
-      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, ..., C)")
+      raise ValueError(
+          "Input must be a tensor with rank >= 3 in layout (N, ..., C)"
+      )
     dims = [0, x.ndim - 1, *range(1, x.ndim - 1)]
     return torch.permute(x, dims)