feature: upgrade neo mxnet to 1.8 (#2235)

Co-authored-by: icywang86rui <[email protected]>

Author: prats13bag

src/sagemaker/image_uri_config/neo-mxnet.json

@@ -2,20 +2,21 @@
     "processors": ["cpu", "gpu"],
     "scope": ["inference"],
     "version_aliases": {
-        "0.12.1": "1.7",
-        "1.0.0": "1.7",
-        "1.1.0": "1.7",
-        "1.2": "1.7",
-        "1.2.0": "1.7",
-        "1.2.1": "1.7",
-        "1.3": "1.7",
-        "1.3.0": "1.7",
-        "1.4": "1.7",
-        "1.4.0": "1.7",
-        "1.4.1": "1.7"
+        "0.12.1": "1.8",
+        "1.0.0": "1.8",
+        "1.1.0": "1.8",
+        "1.2": "1.8",
+        "1.2.0": "1.8",
+        "1.2.1": "1.8",
+        "1.3": "1.8",
+        "1.3.0": "1.8",
+        "1.4": "1.8",
+        "1.4.0": "1.8",
+        "1.4.1": "1.8",
+        "1.7": "1.8"
     },
     "versions": {
-        "1.7": {
+        "1.8": {
             "py_versions": ["py3"],
             "registries": {
                 "af-south-1": "774647643957",

tests/data/mxnet_mnist/mnist_neo.py

@@ -105,7 +105,7 @@ def train(
 
 
 def model_fn(path_to_model_files):
-    import neomxnet  # noqa: F401
+    import neomx  # noqa: F401
 
     ctx = mx.cpu()
     sym, arg_params, aux_params = mx.model.load_checkpoint(
@@ -120,7 +120,7 @@ def model_fn(path_to_model_files):
 
 
 def transform_fn(mod, payload, input_content_type, requested_output_content_type):
-    import neomxnet  # noqa: F401
+    import neomx  # noqa: F401
 
     if input_content_type != "application/vnd+python.numpy+binary":
         raise RuntimeError("Input content type must be application/vnd+python.numpy+binary")

tests/integ/test_neo_mxnet.py

@@ -71,7 +71,7 @@ def test_attach_deploy(
 
         estimator.compile_model(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             output_path=estimator.output_path,
         )
 
@@ -121,7 +121,7 @@ def test_deploy_model(
 
         model.compile(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),
@@ -163,7 +163,7 @@ def test_inferentia_deploy_model(
 
         model.compile(
             target_instance_family=inf_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),