fix lintrunner

Author: ZihangFang

test/end2end/test_temp_allocator_fix.py

@@ -7,19 +7,24 @@
 require temporary memory allocation.
 """
 
-import torch
-import tempfile
 import os
+import tempfile
 from pathlib import Path
-from torch.export import export
-from executorch.exir import EdgeProgramManager, to_edge_transform_and_lower, EdgeCompileConfig
+
+import torch
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
-from executorch.runtime import Verification, Runtime, Program, Method
+from executorch.exir import (
+    EdgeCompileConfig,
+    EdgeProgramManager,
+    to_edge_transform_and_lower,
+)
+from executorch.runtime import Method, Program, Runtime, Verification
+from torch.export import export
 
 
 class TopKModel(torch.nn.Module):
     """Model that uses torch.topk operation which requires temp memory allocation."""
-    
+
     def __init__(self, k=3) -> None:
         super().__init__()
         self.k = k
@@ -32,7 +37,7 @@ def forward(self, x) -> torch.Tensor:
 
 class TopKModelWithOut(torch.nn.Module):
     """Model that uses torch.topk with out parameter which also requires temp memory."""
-    
+
     def __init__(self, k=3) -> None:
         super().__init__()
         self.k = k
@@ -47,42 +52,44 @@ def forward(self, x) -> torch.Tensor:
 def test_topk_without_out_parameter():
     """Test torch.topk without out parameter."""
     print("Testing torch.topk without out parameter...")
-    
+
     model = TopKModel(k=5)
     example_input = (torch.randn(3, 100),)
-    
+
     # Export and compile the model
     with torch.no_grad():
         aten_dialect = export(model, example_input)
-        
+
         backend_dialect = to_edge_transform_and_lower(
             aten_dialect,
             compile_config=EdgeCompileConfig(_check_ir_validity=False),
             partitioner=[XnnpackPartitioner()],
         )
-        
+
         executorch_dialect = backend_dialect.to_executorch()
-        
+
         # Save to temporary file
-        with tempfile.NamedTemporaryFile(suffix='.pte', delete=False) as f:
+        with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f:
             temp_path = f.name
-        
+
         try:
             executorch_dialect.save(temp_path)
-            
+
             # Load and execute with ExecuTorch runtime
             et_runtime = Runtime.get()
             program = et_runtime.load_program(
                 Path(temp_path),
                 verification=Verification.Minimal,
             )
-            
+
             forward = program.load_method("forward")
             outputs = forward.execute(example_input)
-            
-            print(f"✓ Successfully executed topk model: {example_input[0].shape} -> {outputs[0].shape}")
+
+            print(
+                f"✓ Successfully executed topk model: {example_input[0].shape} -> {outputs[0].shape}"
+            )
             return True
-            
+
         finally:
             # Clean up temporary file
             if os.path.exists(temp_path):
@@ -92,42 +99,44 @@ def test_topk_without_out_parameter():
 def test_topk_with_out_parameter():
     """Test torch.topk with out parameter (original failing case)."""
     print("Testing torch.topk with out parameter...")
-    
+
     model = TopKModelWithOut(k=3)
     example_input = (torch.randn(2, 256),)
-    
+
     # Export and compile the model
     with torch.no_grad():
         aten_dialect = export(model, example_input)
-        
+
         backend_dialect = to_edge_transform_and_lower(
             aten_dialect,
             compile_config=EdgeCompileConfig(_check_ir_validity=False),
             partitioner=[XnnpackPartitioner()],
         )
-        
+
         executorch_dialect = backend_dialect.to_executorch()
-        
+
         # Save to temporary file
-        with tempfile.NamedTemporaryFile(suffix='.pte', delete=False) as f:
+        with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f:
             temp_path = f.name
-        
+
         try:
             executorch_dialect.save(temp_path)
-            
+
             # Load and execute with ExecuTorch runtime
             et_runtime = Runtime.get()
             program = et_runtime.load_program(
                 Path(temp_path),
                 verification=Verification.Minimal,
             )
-            
+
             forward = program.load_method("forward")
             outputs = forward.execute(example_input)
-            
-            print(f"✓ Successfully executed topk model with out parameter: {example_input[0].shape} -> {outputs[0].shape}")
+
+            print(
+                f"✓ Successfully executed topk model with out parameter: {example_input[0].shape} -> {outputs[0].shape}"
+            )
             return True
-            
+
         finally:
             # Clean up temporary file
             if os.path.exists(temp_path):
@@ -137,42 +146,44 @@ def test_topk_with_out_parameter():
 def test_larger_topk_operation():
     """Test larger topk operation that would require more temporary memory."""
     print("Testing larger topk operation...")
-    
+
     model = TopKModel(k=50)
     example_input = (torch.randn(5, 1000),)
-    
+
     # Export and compile the model
     with torch.no_grad():
         aten_dialect = export(model, example_input)
-        
+
         backend_dialect = to_edge_transform_and_lower(
             aten_dialect,
             compile_config=EdgeCompileConfig(_check_ir_validity=False),
             partitioner=[XnnpackPartitioner()],
         )
-        
+
         executorch_dialect = backend_dialect.to_executorch()
-        
+
         # Save to temporary file
-        with tempfile.NamedTemporaryFile(suffix='.pte', delete=False) as f:
+        with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f:
             temp_path = f.name
-        
+
         try:
             executorch_dialect.save(temp_path)
-            
+
             # Load and execute with ExecuTorch runtime
             et_runtime = Runtime.get()
             program = et_runtime.load_program(
                 Path(temp_path),
                 verification=Verification.Minimal,
             )
-            
+
             forward = program.load_method("forward")
             outputs = forward.execute(example_input)
-            
-            print(f"✓ Successfully executed large topk model: {example_input[0].shape} -> {outputs[0].shape}")
+
+            print(
+                f"✓ Successfully executed large topk model: {example_input[0].shape} -> {outputs[0].shape}"
+            )
             return True
-            
+
         finally:
             # Clean up temporary file
             if os.path.exists(temp_path):
@@ -183,16 +194,16 @@ def main():
     """Run all tests to verify the temp memory allocation fix."""
     print("Testing temp memory allocation fix for torch.topk operations")
     print("=" * 60)
-    
+
     tests = [
         test_topk_without_out_parameter,
         test_topk_with_out_parameter,
         test_larger_topk_operation,
     ]
-    
+
     passed = 0
     failed = 0
-    
+
     for test in tests:
         try:
             if test():
@@ -202,12 +213,14 @@ def main():
         except Exception as e:
             print(f"✗ Test {test.__name__} failed with exception: {e}")
             failed += 1
-    
+
     print("\n" + "=" * 60)
     print(f"Test Results: {passed} passed, {failed} failed")
-    
+
     if failed == 0:
-        print("✓ All tests passed! The temp memory allocation fix is working correctly.")
+        print(
+            "✓ All tests passed! The temp memory allocation fix is working correctly."
+        )
         return True
     else:
         print("✗ Some tests failed. The fix may not be working correctly.")
@@ -216,4 +229,4 @@ def main():
 
 if __name__ == "__main__":
     success = main()
-    exit(0 if success else 1)
+    exit(0 if success else 1)