fix: fix neo inferentia as compilation target not using framework version

Author: HappyAmazonian

src/sagemaker/model.py

@@ -728,13 +728,22 @@ def _compilation_job_config(
             "Framework": framework.upper(),
         }
 
-        multiple_version_supported_framework_list = ["pytorch", "tensorflow"]
-        if (
-            framework.lower() in multiple_version_supported_framework_list
-            and target_instance_type is not None
-            and re.match("(?=^ml_)(?!ml_inf)", target_instance_type) is not None
-            and framework_version is not None
+        def multi_version_compilation_supported(
+            target_instance_type: str, framework: str, framework_version: str
         ):
+            if target_instance_type and framework and framework_version:
+                framework = framework.lower()
+                multi_version_frameworks_support_mapping = {
+                    "ml_inf": ["pytorch", "tensorflow", "mxnet"],
+                    "ml_ioc": ["pytorch", "tensorflow"],
+                }
+                if re.match("(?=^ml_)", target_instance_type):
+                    return framework in multi_version_frameworks_support_mapping["ml_ioc"]
+                if re.match("(?=^ml_inf)", target_instance_type):
+                    return framework in multi_version_frameworks_support_mapping["ml_inf"]
+            return False
+
+        if multi_version_compilation_supported(target_instance_type, framework, framework_version):
             input_model_config["FrameworkVersion"] = utils.get_short_version(framework_version)
 
         role = self.sagemaker_session.expand_role(role)