lint fixes

Author: PhaneeshB

.github/workflows/ci-amdsharktank.yml

@@ -186,21 +186,21 @@ jobs:
             --iree-hal-target-device=hip \
             --iree-hip-target=gfx942 \
             --iree-device=hip://0 \
-            --ignore=sharktank/tests/layers \
+            --ignore=amdsharktank/tests/layers \
             --device=cuda:0
 
-      - name: Run sharktank layers tests
+      - name: Run amdsharktank layers tests
         if: ${{ !cancelled() }}
         run: |
-          pytest sharktank/tests/layers \
+          pytest amdsharktank/tests/layers \
             --durations=10 \
             --capture=no \
             --log-cli-level=info \
             -v \
             --iree-hal-target-device=hip \
             --iree-hip-target=gfx942 \
             --iree-device=hip://0 \
-            --parameters=/shark-dev/ossci-models/llama_3_1/405b/fp4/fp4_preshuffled_2025_09_12.irpa \
+            --parameters=/amdshark-dev/ossci-models/llama_3_1/405b/fp4/fp4_preshuffled_2025_09_12.irpa \
             --device=cuda:0 \
             -n 4
 

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,