lint fixes

Author: PhaneeshB

amdsharktank/amdsharktank/utils/export.py

@@ -327,4 +327,4 @@ def _(module, *fn_args):
     export_output = aot.export(fxb, import_symbolic_shape_expressions=True)
     export_output.save_mlir(mlir_path)
 
-    return export_output
+    return export_output

amdsharktank/amdsharktank/utils/iree.py

@@ -30,7 +30,11 @@
     torch_tree_flatten,
 )
 from amdsharktank.utils import verify_exactly_one_is_not_none
-from amdsharktank.utils.export import export_torch_module_to_mlir_file, get_torch_eager_output, _as_tuple
+from amdsharktank.utils.export import (
+    export_torch_module_to_mlir_file,
+    get_torch_eager_output,
+    _as_tuple,
+)
 from .tree import Tree
 from iree.runtime import FileHandle
 import iree.runtime

amdsharktank/tests/layers/rotary_embedding_hf_test.py

@@ -26,7 +26,11 @@
 )
 from amdsharktank import ops
 from amdsharktank.utils.logging import get_logger
-from amdsharktank.utils.testing import TempDirTestBase, assert_tensor_close, is_hip_condition
+from amdsharktank.utils.testing import (
+    TempDirTestBase,
+    assert_tensor_close,
+    is_hip_condition,
+)
 import iree.compiler
 from iree.turbine.aot import (
     FxProgramsBuilder,