shark migration update

Author: PhaneeshB

amdsharktank/amdsharktank/utils/export.py

@@ -6,6 +6,7 @@
 
 from typing import Callable, Optional, Any
 from os import PathLike
+from pathlib import Path
 import functools
 
 import torch
@@ -16,7 +17,13 @@
 from torch.utils._pytree import tree_structure, tree_unflatten, tree_flatten
 from amdsharktank.types.tensors import ShardedTensor
 from amdsharktank.types.theta import mark_export_external_theta
-from amdsharktank.layers import BaseLayer, ThetaLayer
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from amdsharktank.layers import BaseLayer, ThetaLayer
+
+# from amdsharktank.layers import BaseLayer, ThetaLayer
 
 
 def flatten_signature(
@@ -180,7 +187,7 @@ def flat_fn(*args, **kwargs):
 
 
 def export_model_mlir(
-    model: BaseLayer,
+    model: "BaseLayer",
     output_path: PathLike,
     *,
     function_batch_sizes_map: Optional[dict[Optional[str], list[int]]] = None,
@@ -202,7 +209,7 @@ def export_model_mlir(
 
     assert not (function_batch_sizes_map is not None and batch_sizes is not None)
 
-    if isinstance(model, ThetaLayer):
+    if isinstance(model, "ThetaLayer"):
         mark_export_external_theta(model.theta)
 
     if batch_sizes is not None:
@@ -317,7 +324,7 @@ def export_torch_module_to_mlir_file(
     def _(module, *fn_args):
         return module.forward(*fn_args)
 
-    export_output = export(fxb, import_symbolic_shape_expressions=True)
+    export_output = aot.export(fxb, import_symbolic_shape_expressions=True)
     export_output.save_mlir(mlir_path)
 
     return export_output

amdsharktank/amdsharktank/utils/iree.py

@@ -35,9 +35,6 @@
 from iree.runtime import FileHandle
 import iree.runtime
 
-from iree.turbine import aot
-from iree.turbine.aot import export
-
 
 if TYPE_CHECKING:
     from ..layers import ModelConfig

amdsharktank/tests/layers/ffn_with_iree_test.py

@@ -5,9 +5,9 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 import torch
 import pytest
-from sharktank.utils.iree import run_iree_vs_torch_eager
-from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-from sharktank.utils.testing import is_hip_condition
+from amdsharktank.utils.iree import run_iree_vs_torch_eager
+from amdsharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from amdsharktank.utils.testing import is_hip_condition
 
 
 class FFN(torch.nn.Module):

amdsharktank/tests/layers/linear_with_iree_test.py

@@ -7,9 +7,9 @@
 import torch
 import pytest
 import gc
-from sharktank.utils.iree import run_iree_vs_torch_eager
-from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-from sharktank.utils.testing import is_hip_condition
+from amdsharktank.utils.iree import run_iree_vs_torch_eager
+from amdsharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from amdsharktank.utils.testing import is_hip_condition
 
 
 class Linear(torch.nn.Module):

amdsharktank/tests/layers/output_lm_head_with_iree_test.py

@@ -7,14 +7,14 @@
 import torch
 import pytest
 from pathlib import Path
-from sharktank.utils.iree import (
+from amdsharktank.utils.iree import (
     run_iree_vs_torch_eager,
 )
-from sharktank.layers import LinearLayer, RMSNormLayer
-from sharktank.types import Dataset, Theta
-from sharktank.layers.configs import LlamaModelConfig
-from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-from sharktank.utils.testing import is_hip_condition, validate_and_get_irpa_path
+from amdsharktank.layers import LinearLayer, RMSNormLayer
+from amdsharktank.types import Dataset, Theta
+from amdsharktank.layers.configs import LlamaModelConfig
+from amdsharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from amdsharktank.utils.testing import is_hip_condition, validate_and_get_irpa_path
 
 
 class OutputLMHead(torch.nn.Module):
@@ -136,7 +136,7 @@ def test_output_lm_head_mock():
     torch.manual_seed(42)
 
     # Mock configuration - provide all required parameters
-    from sharktank.layers.configs import LlamaHParams
+    from amdsharktank.layers.configs import LlamaHParams
 
     # Create LlamaHParams with all required parameters
     hp = LlamaHParams(
@@ -160,7 +160,7 @@ def test_output_lm_head_mock():
     )
 
     # Create mock theta with synthetic weights
-    from sharktank.types import DefaultPrimitiveTensor
+    from amdsharktank.types import DefaultPrimitiveTensor
 
     # Mock output_norm weights
     output_norm_weight = torch.randn(hp.embedding_length, dtype=torch.float32)

amdsharktank/tests/layers/rms_norm_with_iree_test.py

@@ -7,9 +7,9 @@
 import torch
 import pytest
 import gc
-from sharktank.utils.iree import run_iree_vs_torch_eager
-from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-from sharktank.utils.testing import is_hip_condition
+from amdsharktank.utils.iree import run_iree_vs_torch_eager
+from amdsharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from amdsharktank.utils.testing import is_hip_condition
 
 
 class RMSNorm(torch.nn.Module):

amdsharktank/tests/layers/token_embedding_with_iree_test.py

@@ -7,13 +7,13 @@
 import torch
 import pytest
 from pathlib import Path
-from sharktank.layers.token_embedding import TokenEmbeddingLayer
-from sharktank.types.theta import Dataset
-from sharktank.utils.iree import (
+from amdsharktank.layers.token_embedding import TokenEmbeddingLayer
+from amdsharktank.types.theta import Dataset
+from amdsharktank.utils.iree import (
     run_iree_vs_torch_eager,
 )
-from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-from sharktank.utils.testing import is_hip_condition, validate_and_get_irpa_path
+from amdsharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from amdsharktank.utils.testing import is_hip_condition, validate_and_get_irpa_path
 
 
 class TokenEmbeddingSmall(torch.nn.Module):