# Copyright (c) 2021, Apple Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-3-clause license that can be
# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import copy
from io import BytesIO
import itertools
import numpy as np
import os
from os import getcwd, chdir
from os.path import exists
from PIL import Image
import pytest
from shutil import rmtree
from tempfile import mkdtemp
import urllib.request

import coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.frontend.torch.test.testing_utils import _copy_input_data
from coremltools.converters.mil.mil import get_new_symbol, Program, Function
from coremltools.converters.mil.testing_utils import get_op_types_in_program
from coremltools.proto import FeatureTypes_pb2 as ft

from coremltools._deps import (
    _HAS_TF_1,
    _HAS_TF_2,
    _HAS_TORCH,
    MSG_TF1_NOT_FOUND,
    MSG_TF2_NOT_FOUND,
    MSG_TORCH_NOT_FOUND,
)

if _HAS_TORCH:
    import torch
    import torchvision

if _HAS_TF_1 or _HAS_TF_2:
    import tensorflow as tf

if _HAS_TF_2:
    import requests
    from tensorflow import keras
    from tensorflow.keras import layers


###############################################################################
# Note: all tests are also used as examples such as in readme.md as a reference
# Whenever any of the following test fails, we should update API documentations
# Each test case is expected to be runnable and self-complete, then sync to the
# documentation pages as API example code snippet.
###############################################################################


@pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND)
@pytest.mark.skipif(ct.utils._macos_version() < (10, 15), reason='Model produces specification 4.')
class TestTensorFlow1ConverterExamples:

    @staticmethod
    def test_convert_from_frozen_graph(tmpdir):
        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input")
            y = tf.nn.relu(x, name="output")

        mlmodel = ct.convert(graph, compute_units=ct.ComputeUnit.CPU_ONLY)

        test_input = np.random.rand(1, 2, 3) - 0.5
        with tf.compat.v1.Session(graph=graph) as sess:
            expected_val = sess.run(y, feed_dict={x: test_input})
        results = mlmodel.predict({"input": test_input})
        np.testing.assert_allclose(results["output"], expected_val)

    @staticmethod
    def test_convert_from_frozen_graph_file(tmpdir):
        # create the model to convert

        # write a toy frozen graph
        # Note that we usually needs to run freeze_graph() on tf.Graph()
        # skipping here as this toy model does not contain any variables
        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input")
            y = tf.nn.relu(x, name="output")

        save_path = str(tmpdir)
        tf.io.write_graph(graph, save_path, "frozen_graph.pb", as_text=False)

        # Create a test sample
        # -0.5 to have some negative values
        test_input = np.random.rand(1, 2, 3) - 0.5
        with tf.compat.v1.Session(graph=graph) as sess:
            expected_val = sess.run(y, feed_dict={x: test_input})

        # The input `.pb` file is a frozen graph format that usually
        # generated by TensorFlow's utility function `freeze_graph()`
        pb_path = os.path.join(save_path, "frozen_graph.pb")

        # 3 ways to specify inputs:
        # (1) Fully specify inputs
        mlmodel = ct.convert(
            pb_path,
            # We specify inputs with name matching the placeholder name.
            inputs=[ct.TensorType(name="input", shape=(1, 2, 3))],
            outputs=["output"],
        )

        # (2) Specify input TensorType without name (when there's only one
        # input)
        mlmodel = ct.convert(
            pb_path,
            # TensorType name is optional when there's only one input.
            inputs=[ct.TensorType(shape=(1, 2, 3))],
            outputs=["output"],
        )

        # (3) Not specify inputs at all. `inputs` is optional for TF. When
        # inputs is not specified, convert() infers inputs from Placeholder
        # nodes.
        mlmodel = ct.convert(pb_path, outputs=["output"], compute_units=ct.ComputeUnit.CPU_ONLY)

        results = mlmodel.predict({"input": test_input})
        np.testing.assert_allclose(results["output"], expected_val)
        mlmodel_path = os.path.join(save_path, "model.mlmodel")
        # Save the converted model
        mlmodel.save(mlmodel_path)

        results = mlmodel.predict({"input": test_input})
        np.testing.assert_allclose(results["output"], expected_val, atol=1e-3)

    @staticmethod
    def test_convert_from_saved_model_dir(tmpdir):
        # Sample input
        test_input = np.random.rand(1, 3, 5) - 0.5

        # create the model to convert
        with tf.compat.v1.Session() as sess:
            x = tf.placeholder(shape=(1, 3, 5), dtype=tf.float32)
            y = tf.nn.relu(x)

            expected_val = sess.run(y, feed_dict={x: test_input})

        # Save model as SavedModel
        inputs = {"x": x}
        outputs = {"y": y}
        save_path = str(tmpdir)
        tf.compat.v1.saved_model.simple_save(sess, save_path, inputs, outputs)

        # SavedModel directory generated by TensorFlow 1.x
        # when converting from SavedModel dir, inputs / outputs are optional
        mlmodel = ct.convert(save_path, compute_units=ct.ComputeUnit.CPU_ONLY)

        # Need input output names to call mlmodel
        # x.name == 'Placeholder:0'. Strip out ':0'
        input_name = x.name.split(":")[0]
        results = mlmodel.predict({input_name: test_input})
        # y.name == 'Relu:0'. output_name == 'Relu'
        output_name = y.name.split(":")[0]
        np.testing.assert_allclose(results[output_name], expected_val)


    @staticmethod
    def test_freeze_and_convert_matmul_graph():
        # testing : https://coremltools.readme.io/docs/tensorflow-1#export-as-frozen-graph-and-convert
        graph = tf.Graph()
        with graph.as_default():
            x = tf.placeholder(tf.float32, shape=[None, 20], name="input")
            W = tf.Variable(tf.truncated_normal([20, 10], stddev=0.1))
            b = tf.Variable(tf.ones([10]))
            y = tf.matmul(x, W) + b
            output_names = [y.op.name]

        from tensorflow.python.tools.freeze_graph import freeze_graph

        model_dir = mkdtemp()
        graph_def_file = os.path.join(model_dir, 'tf_graph.pb')
        checkpoint_file = os.path.join(model_dir, 'tf_model.ckpt')
        frozen_graph_file = os.path.join(model_dir, 'tf_frozen.pb')

        with tf.Session(graph=graph) as sess:
            # initialize variables
            sess.run(tf.global_variables_initializer())
            # save graph definition somewhere
            tf.train.write_graph(sess.graph, model_dir, graph_def_file, as_text=False)
            # save the weights
            saver = tf.train.Saver()
            saver.save(sess, checkpoint_file)

            # take the graph definition and weights
            # and freeze into a single .pb frozen graph file
            freeze_graph(input_graph=graph_def_file,
                         input_saver="",
                         input_binary=True,
                         input_checkpoint=checkpoint_file,
                         output_node_names=",".join(output_names),
                         restore_op_name="save/restore_all",
                         filename_tensor_name="save/Const:0",
                         output_graph=frozen_graph_file,
                         clear_devices=True,
                         initializer_nodes="")
        print("Tensorflow frozen graph saved at {}".format(frozen_graph_file))

        mlmodel = ct.convert(frozen_graph_file)
        # optionally, you can save model to disk
        # mlmodel.save(frozen_graph_file.replace("pb", "mlmodel"))
        try:
            rmtree(model_dir)
        except:
            pass


@pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
class TestTensorFlow2ConverterExamples:
    def setup_class(self):
        self._cwd = getcwd()
        self._temp_dir = mkdtemp()
        # step into temp directory as working directory
        # to make the user-facing examples cleaner
        chdir(self._temp_dir)

        # create toy models for conversion examples
        # write a toy tf.keras HDF5 model
        tf_keras_model = tf.keras.Sequential(
            [
                tf.keras.layers.Flatten(input_shape=(28, 28)),
                tf.keras.layers.Dense(128, activation=tf.nn.relu),
                tf.keras.layers.Dense(10, activation=tf.nn.softmax),
            ]
        )
        tf_keras_model.save("./tf_keras_model.h5")

        # write a toy SavedModel directory
        tf_keras_model.save("./saved_model", save_format="tf")

    def teardown_class(self):
        chdir(self._cwd)
        if exists(self._temp_dir):
            rmtree(self._temp_dir)

    @staticmethod
    def test_convert_tf_keras_h5_file(tmpdir):
        for file_extension in ("h5", ".hdf5"):
            x = tf.keras.Input(shape=(32,), name="input")
            y = tf.keras.layers.Dense(16, activation="softmax")(x)
            keras_model = tf.keras.Model(x, y)
            save_dir = str(tmpdir)
            path = os.path.join(save_dir, "tf_keras_model." + file_extension)
            keras_model.save(path)

            mlmodel = ct.convert(path)

            test_input = np.random.rand(2, 32)
            expected_val = keras_model(test_input)
            results = mlmodel.predict({"input": test_input})
            np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4)

    @staticmethod
    def test_convert_tf_keras_model():
        x = tf.keras.Input(shape=(32,), name="input")
        y = tf.keras.layers.Dense(16, activation="softmax")(x)
        keras_model = tf.keras.Model(x, y)

        mlmodel = ct.convert(keras_model)

        test_input = np.random.rand(2, 32)
        expected_val = keras_model(test_input)
        results = mlmodel.predict({"input": test_input})
        np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4)

    @staticmethod
    @pytest.mark.parametrize(
            "dtype", ['default', 'mil_type', 'np type'])
    def test_convert_tf_keras_applications_model(dtype):
        tf_keras_model = tf.keras.applications.MobileNet(
            weights="imagenet", input_shape=(224, 224, 3)
        )

        # inputs / outputs are optional, we can get from tf.keras model
        # this can be extremely helpful when we want to extract sub-graphs
        input_name = tf_keras_model.inputs[0].name.split(":")[0]
        # note that the `convert()` requires tf.Graph's outputs instead of
        # tf.keras.Model's outputs, to access that, we can do the following
        output_name = tf_keras_model.outputs[0].name.split(":")[0]
        tf_graph_output_name = output_name.split("/")[-1]

        if dtype == 'default':
            dtype = None
        elif dtype == 'mil_type':
            dtype = ct.converters.mil.types.fp32
        else:
            dtype = np.float32

        mlmodel = ct.convert(
            tf_keras_model,
            inputs=[ct.TensorType(shape=(1, 224, 224, 3), dtype=dtype)],
            outputs=[tf_graph_output_name],
        )
        mlmodel.save("./mobilenet.mlmodel")

    @staticmethod
    def test_convert_from_saved_model_dir():
        # SavedModel directory generated by TensorFlow 2.x
        mlmodel = ct.convert("./saved_model")
        mlmodel.save("./model.mlmodel")


    @staticmethod
    def test_keras_custom_layer_model():
        # testing : https://coremltools.readme.io/docs/tensorflow-2#conversion-from-user-defined-models

        class CustomDense(layers.Layer):
            def __init__(self, units=32):
                super(CustomDense, self).__init__()
                self.units = units

            def build(self, input_shape):
                self.w = self.add_weight(
                    shape=(input_shape[-1], self.units),
                    initializer="random_normal",
                    trainable=True,
                )
                self.b = self.add_weight(
                    shape=(self.units,), initializer="random_normal", trainable=True
                )

            def call(self, inputs):
                return tf.matmul(inputs, self.w) + self.b

        inputs = keras.Input((4,))
        outputs = CustomDense(10)(inputs)
        model = keras.Model(inputs, outputs)
        ct.convert(model)

    @staticmethod
    def test_concrete_function_conversion():
        # testing : https://coremltools.readme.io/docs/tensorflow-2#conversion-from-user-defined-models

        @tf.function(input_signature=[tf.TensorSpec(shape=(6,), dtype=tf.float32)])
        def gelu_tanh_activation(x):
            a = (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))
            y = 0.5 * (1.0 + tf.tanh(a))
            return x * y

        conc_func = gelu_tanh_activation.get_concrete_function()
        mlmodel = ct.convert([conc_func])


    @staticmethod
    def test_quickstart_example():
        # testing: https://coremltools.readme.io/docs/introductory-quickstart#quickstart-example

        # Download MobileNetv2 (using tf.keras)
        keras_model = tf.keras.applications.MobileNetV2(
            weights="imagenet",
            input_shape=(224, 224, 3,),
            classes=1000,
        )

        # Download class labels (from a separate file)
        label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
        class_labels = urllib.request.urlopen(label_url).read().splitlines()
        class_labels = class_labels[1:]  # remove the first class which is background
        assert len(class_labels) == 1000

        # make sure entries of class_labels are strings
        for i, label in enumerate(class_labels):
            if isinstance(label, bytes):
                class_labels[i] = label.decode("utf8")

        image_input = ct.ImageType(shape=(1, 224, 224, 3,),
                                   bias=[-1, -1, -1], scale=1 / 127)

        # set class labels
        classifier_config = ct.ClassifierConfig(class_labels)

        # Convert the model using the Unified Conversion API
        model = ct.convert(
            keras_model, inputs=[image_input], classifier_config=classifier_config,
        )

        # Set feature descriptions (these show up as comments in XCode)
        input_name = model.input_description._fd_spec[0].name
        model.input_description[input_name] = "Input image to be classified"
        model.output_description["classLabel"] = "Most likely image category"
        model.author = "Original Paper: Mark Sandler, Andrew Howard, "\
                        "Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen"
        model.license = "Please see https://github.com/tensorflow/tensorflow "\
                        "for license information, and "\
                        "https://github.com/tensorflow/models/tree/master/research/slim/nets/mobilenet"\
                        "for the original source of the model."
        model.short_description = "Detects the dominant objects present in an"\
                                 "image from a set of 1001 categories such as trees, animals,"\
                                 "food, vehicles, person etc. The top-1 accuracy"\
                                " from the original publication is 74.7%."
        model.version = "2.0"

        # get an image
        img_url = 'https://files.readme.io/02e3586-daisy.jpg'
        response = requests.get(img_url)
        img = Image.open(BytesIO(response.content))

        # Use PIL to load and resize the image to expected size
        example_image = img.resize((224, 224))

        # Make a prediction using Core ML
        out_dict = model.predict({input_name: example_image})

        # Print out top-1 prediction
        assert out_dict["classLabel"] == "daisy"

@pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
class TestPyTorchConverterExamples:
    @staticmethod
    def test_convert_torch_vision_mobilenet_v2(tmpdir):
        """
        In this example, we'll instantiate a PyTorch classification model and convert
        it to Core ML.
        """

        """
        Here we instantiate our model. In a real use case this would be your trained
        model.
        """
        model = torchvision.models.mobilenet_v2()

        """
        The next thing we need to do is generate TorchScript for the model. The easiest
        way to do this is by tracing it.
        """

        """
        It's important that a model be in evaluation mode (not training mode) when it's
        traced. This makes sure things like dropout are disabled.
        """
        model.eval()

        """
        Tracing takes an example input and traces its flow through the model. Here we
        are creating an example image input.

        The rank and shape of the tensor will depend on your model use case. If your
        model expects a fixed size input, use that size here. If it can accept a
        variety of input sizes, it's generally best to keep the example input small to
        shorten how long it takes to run a forward pass of your model. In all cases,
        the rank of the tensor must be fixed.
        """
        example_input = torch.rand(1, 3, 256, 256)

        """
        Now we actually trace the model. This will produce the TorchScript that the
        CoreML converter needs.
        """
        traced_model = torch.jit.trace(model, example_input)

        """
        Now with a TorchScript representation of the model, we can call the CoreML
        converter. The converter also needs a description of the input to the model,
        where we can give it a convenient name.
        """
        mlmodel = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
        )

        """
        Now with a conversion complete, we can save the MLModel and run inference.
        """
        save_path = os.path.join(str(tmpdir), "mobilenet_v2.mlmodel")
        mlmodel.save(save_path)

        """
        Running predict() is only supported on macOS.
        """
        if ct.utils._is_macos():
            results = mlmodel.predict({"input": example_input.numpy()})
            assert isinstance(results, dict)
            
    @staticmethod
    def test_int64_inputs():

        num_tokens = 3
        embedding_size = 5

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.embedding = torch.nn.Embedding(num_tokens,
                    embedding_size)

            def forward(self, x):
                return self.embedding(x)

        model = TestModule()
        model.eval()

        example_input = torch.randint(high=num_tokens, size=(2,), dtype=torch.int64)
        traced_model = torch.jit.trace(model, example_input)
        mlmodel = ct.convert(
            traced_model,
            inputs=[
                ct.TensorType(
                    name="input",
                    shape=example_input.shape,
                    dtype=example_input.numpy().dtype,
                )
            ],
        )

        # running predict() is supported on macOS
        if ct.utils._is_macos():
            result = mlmodel.predict(
                {"input": example_input.detach().numpy().astype(np.float32)}
            )

            # Verify outputs
            expected = model(example_input)
            name = list(result.keys())[0]
            np.testing.assert_allclose(result[name], expected.detach().numpy())

        # Duplicated inputs are invalid
        with pytest.raises(ValueError, match=r"Duplicated inputs"):
            mlmodel = ct.convert(
                traced_model,
                inputs=[
                    ct.TensorType(
                        name="input",
                        shape=example_input.shape,
                        dtype=example_input.numpy().dtype,
                    ),
                    ct.TensorType(
                        name="input",
                        shape=example_input.shape,
                        dtype=example_input.numpy().dtype,
                    ),
                ],
            )

        # Outputs must be of type ct.ImageType or ct.TensorType
        with pytest.raises(ValueError, match=r"must be a list of type ct.TensorType or ct.ImageType"):
            mlmodel = ct.convert(
                traced_model,
                inputs=[
                    ct.TensorType(
                        name="input",
                        shape=example_input.shape,
                        dtype=example_input.numpy().dtype,
                    ),
                ],
                outputs=["output"],
            )

    @staticmethod
    def test_fully_dynamic_inputs():
        """
        All dims of the inputs are dynamic, and write to slice to one of the
        inputs.
        """

        class Model(torch.nn.Module):
            def __init__(self, index):
                super(Model, self).__init__()
                self.index = index

            def forward(self, x, y):
                x[:, int(self.index.item())] = 0.0
                y = y.unsqueeze(0)
                return y, x

        model = Model(torch.tensor(3))
        scripted_model = torch.jit.script(model)

        mlmodel = ct.convert(
            scripted_model,
            inputs=[
                ct.TensorType("x", shape=(ct.RangeDim(), ct.RangeDim())),
                ct.TensorType("y", shape=(ct.RangeDim(), ct.RangeDim()))
            ],
        )

        # running predict() is supported on macOS
        if ct.utils._is_macos():
            x, y = torch.rand(2, 4), torch.rand(1, 2)
            torch_input = _copy_input_data([x, y])
            torch_res = model(*torch_input)
            results = mlmodel.predict({"x": x.cpu().detach().numpy(),
              "y": y.cpu().detach().numpy()})
            for i, name in enumerate(mlmodel.output_description):
                np.testing.assert_allclose(torch_res[i], results[name])

            x, y = torch.rand(1, 6), torch.rand(2, 3)
            torch_input = _copy_input_data([x, y])
            torch_res = model(*torch_input)
            results = mlmodel.predict({"x": x.cpu().detach().numpy(),
              "y": y.cpu().detach().numpy()})
            for i, name in enumerate(mlmodel.output_description):
                np.testing.assert_allclose(torch_res[i], results[name])


class TestMILExamples:
    @staticmethod
    def test_tutorial():
        @mb.program(
            input_specs=[mb.TensorSpec(shape=(1, 100, 100, 3))]
        )
        def prog(x):
            x = mb.relu(x=x, name="relu")
            x = mb.transpose(x=x, perm=[0, 3, 1, 2], name="transpose")
            x = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=False, name="reduce")
            x = mb.log(x=x, name="log")
            y = mb.add(x=1, y=2)
            return x

        print("prog:\n", prog)

        # Convert and verify
        mlmodel = ct.convert(prog)

        # running predict() is only supported on macOS
        if ct.utils._is_macos():
            prediction = mlmodel.predict(
                {"x": np.random.rand(1, 100, 100, 3).astype(np.float32)}
            )
            assert len(prediction) == 1


class TestInvalidInput:
    @staticmethod
    def test_rank0_inputs_mil():
        with pytest.raises(ValueError, match=r"Rank-0"):
            @mb.program(
                input_specs=[mb.TensorSpec(shape=()),]
            )
            def prog(x):
                return x

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_rank0_inputs_torch():
        """Similar to TestPyTorchConverterExamples::test_int64_inputs but
        using rank-0 int input.
        """

        num_tokens = 3
        embedding_size = 5

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.embedding = torch.nn.Embedding(num_tokens,
                    embedding_size)

            def forward(self, x):
                return self.embedding(x)

        model = TestModule()
        model.eval()

        example_input = torch.tensor(1)
        traced_model = torch.jit.trace(model, example_input)
        with pytest.raises(ValueError, match=r"Rank-0"):
            mlmodel = ct.convert(
                traced_model,
                inputs=[
                    ct.TensorType(
                        name="input",
                        shape=example_input.shape,
                        dtype=example_input.numpy().dtype,
                    )
                ],
            )

@pytest.mark.skipif(ct.utils._macos_version() < (10, 15), reason='Model produces specification 4.')
class TestInputs:
    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND)
    def test_input_noname():
        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input")
            x1 = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input_1")
            y = tf.nn.relu(x, name="output")
            y1 = tf.nn.relu(x1, name="output_1")

        with pytest.raises(ValueError) as e:
            model = ct.convert(
                graph,
                inputs=[ct.TensorType(shape=(1, 2, 3))]
            )
        expected_error = "Multiple inputs are found in graph, but no input name was provided"
        assert expected_error == str(e.value)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND)
    def test_input_wrongname():
        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input")
            x1 = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input_1")
            y = tf.nn.relu(x, name="output")
            y1 = tf.nn.relu(x1, name="output_1")

        with pytest.raises(ValueError) as e:
            model = ct.convert(
                graph,
                inputs=[ct.TensorType(shape=(1, 2, 3), name="wrong_input")]
            )
        expected_error = "Multiple inputs are found in graph, but no input name was provided"
        expected_error = "Input ({}) provided is not found in given tensorflow graph. Placeholders in graph are: {}".format("wrong_input", ["input", "input_1"])
        assert expected_error == str(e.value)

    @staticmethod
    @pytest.mark.skipif(not ct.utils._is_macos(), reason="Platform is not Mac OS")
    def test_unsanitized_input_name_during_prediction():
        '''
        input name : "x/0" becomes "x_0" due to name sanitization applied during conversion
        '''
        prog = Program()
        func_inputs = {"x/0": mb.placeholder(shape=[2, 3]),
                       "y": mb.placeholder(shape=[2, 3])}
        with Function(func_inputs) as ssa_fun:
            x, y = ssa_fun.inputs["x/0"], ssa_fun.inputs["y"]
            x = mb.relu(x=x, name="relu")
            z = mb.add(x=x, y=y, name="out")
            ssa_fun.set_outputs([z])
        prog.add_function("main", ssa_fun)

        mlmodel = ct.convert(prog)

        with pytest.raises(KeyError) as error_info:
            mlmodel.predict(
                {"x/0": np.random.rand(2, 3).astype(np.float32),
                 "y": np.random.rand(2, 3).astype(np.float32)}
            )
        error_str = str(error_info.value)
        assert "does not match any of the model input" in error_str

    @staticmethod
    def _test_variant_input_type_prediction(to_tensor):
        prog = Program()
        func_inputs = {"x": mb.placeholder(shape=[2, 3]),
                       "y": mb.placeholder(shape=[2, 3])}
        with Function(func_inputs) as ssa_fun:
            x, y = ssa_fun.inputs["x"], ssa_fun.inputs["y"]
            x = mb.relu(x=x, name="relu")
            z = mb.add(x=x, y=y, name="out")
            ssa_fun.set_outputs([z])
        prog.add_function("main", ssa_fun)

        mlmodel = ct.convert(prog)
        x_numpy = np.random.rand(2, 3)
        y_numpy = np.random.rand(2, 3)
        out_by_numpy = mlmodel.predict(
            {"x": x_numpy,
             "y": y_numpy}
        )
        out_by_tensor = mlmodel.predict(
            {"x": to_tensor(x_numpy),
             "y": to_tensor(y_numpy)}
        )
        np.allclose(out_by_numpy["out"], out_by_tensor["out"])

    @staticmethod
    @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions")
    def test_list_predict_input():
        TestInputs._test_variant_input_type_prediction(lambda x: x.tolist())

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_1, reason="requires TensorFlow")
    @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions")
    def test_tf_predict_input():
        TestInputs._test_variant_input_type_prediction(tf.convert_to_tensor)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason="requires Torch")
    @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions")
    def test_torch_predict_input():
        TestInputs._test_variant_input_type_prediction(torch.tensor)


class TestFlexibleShape:
    @staticmethod
    @pytest.mark.parametrize(
            "use_symbol", [True, False])
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2keras_shared_range_dim(use_symbol):
        # Test examples in https://coremltools.readme.io/docs/flexible-inputs
        input_dim = 3
        # None denotes seq_len dimension
        x1 = tf.keras.Input(shape=(None,input_dim), name="seq1")
        x2 = tf.keras.Input(shape=(None,input_dim), name="seq2")
        y = x1 + x2
        keras_model = tf.keras.Model(inputs=[x1, x2], outputs=[y])

        # One RangeDim shared by two inputs
        if use_symbol:
            seq_len_dim = ct.RangeDim(symbol='seq_len')
        else:
            # symbol is optional
            seq_len_dim = ct.RangeDim()
        seq1_input = ct.TensorType(name="seq1", shape=(1, seq_len_dim, input_dim))
        seq2_input = ct.TensorType(name="seq2", shape=(1, seq_len_dim, input_dim))
        mlmodel = ct.convert(keras_model,
                inputs=[seq1_input, seq2_input])

        batch = 1
        seq_len = 5
        test_input_x1 = np.random.rand(batch, seq_len, input_dim).astype(np.float32)
        test_input_x2 = np.random.rand(batch, seq_len, input_dim).astype(np.float32)
        expected_val = keras_model([test_input_x1, test_input_x2])
        if ct.utils._is_macos():
            results = mlmodel.predict({
                "seq1": test_input_x1,
                "seq2": test_input_x2})
            np.testing.assert_allclose(results["Identity"], expected_val,
                rtol=1e-4, atol=1e-3)


    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2keras_incorrect_range_dim():
        import tensorflow as tf

        input_dim = 3
        # None denotes seq_len dimension
        x1 = tf.keras.Input(shape=(None,input_dim), name="seq1")
        y = x1 + 1
        keras_model = tf.keras.Model(inputs=[x1], outputs=[y])

        # Incorrectly using -1 instead of ct.RangeDim
        # One RangeDim shared by two inputs
        with pytest.raises(ValueError,
            match=r"Can\'t convert to CoreML shaping"):
            seq1_input = ct.TensorType(name="seq1", shape=(1, -1, input_dim))
            mlmodel = ct.convert(keras_model, inputs=[seq1_input])

    @staticmethod
    @pytest.mark.parametrize(
            "use_symbol", [True, False])
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2keras_outofbound_range_dim(use_symbol):
        # Test examples in https://coremltools.readme.io/docs/flexible-inputs
        import tensorflow as tf

        input_dim = 3
        # None denotes seq_len dimension
        x = tf.keras.Input(shape=(None,input_dim), name="seq")
        y = x * 2
        keras_model = tf.keras.Model(inputs=[x], outputs=[y])

        if use_symbol:
            seq_len_dim = ct.RangeDim(symbol='sequence_len', lower_bound=3,
                    upper_bound=5)
        else:
            seq_len_dim = ct.RangeDim(lower_bound=3, upper_bound=5)
        seq_input = ct.TensorType(name="seq", shape=(1, seq_len_dim, input_dim))
        mlmodel = ct.convert(keras_model, inputs=[seq_input])

        # seq_len is within bound
        batch = 1
        seq_len = 3
        test_input_x = np.random.rand(batch, seq_len, input_dim).astype(np.float32)
        expected_val = keras_model([test_input_x])
        if ct.utils._is_macos():
            results = mlmodel.predict({"seq": test_input_x})
            np.testing.assert_allclose(results["Identity"], expected_val,
                rtol=1e-4, atol=1e-3)

            # seq_len below/above lower_bound/upper_bound
            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                seq_len = 2
                test_input_x = np.random.rand(batch, seq_len,
                        input_dim).astype(np.float32)
                results = mlmodel.predict({"seq": test_input_x})

            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                seq_len = 6
                test_input_x = np.random.rand(batch, seq_len,
                        input_dim).astype(np.float32)
                results = mlmodel.predict({"seq": test_input_x})


    @staticmethod
    @pytest.mark.parametrize(
            "use_symbol", [True, False])
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_range_dim(use_symbol):

        num_tokens = 3
        embedding_size = 5

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.embedding = torch.nn.Embedding(num_tokens, embedding_size)

            def forward(self, x):
                return self.embedding(x)

        model = TestModule()
        model.eval()

        example_input = torch.randint(high=num_tokens, size=(2,),
                dtype=torch.int64)
        traced_model = torch.jit.trace(model, example_input)

        if use_symbol:
            seq_len_dim = ct.RangeDim(symbol='seq_length')
        else:
            # symbol is optional
            seq_len_dim = ct.RangeDim()
        seq_input = ct.TensorType(name="input", shape=(seq_len_dim,),
                dtype=np.int64)
        mlmodel = ct.convert(
            traced_model,
            inputs=[seq_input],
        )

        if ct.utils._is_macos():
            result = mlmodel.predict(
                {"input": example_input.detach().numpy().astype(np.float32)}
            )

            # Verify outputs
            expected = model(example_input)
            name = list(result.keys())[0]
            np.testing.assert_allclose(result[name], expected.detach().numpy())

            # Try example of different length
            example_input2 = torch.randint(high=num_tokens, size=(99,),
                    dtype=torch.int64)
            result = mlmodel.predict(
                {"input": example_input2.detach().numpy().astype(np.float32)}
            )
            expected = model(example_input2)
            name = list(result.keys())[0]
            np.testing.assert_allclose(result[name], expected.detach().numpy())

    @staticmethod
    @pytest.mark.parametrize(
            "variable_length", [True, False])
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_range_dim_lstm(variable_length):
        """
        This example shows how to run LSTM with previous hidden / cell states
        """

        input_size = 3
        hidden_size = 2

        class TestNet(torch.nn.Module):
            def __init__(self):
              super(TestNet, self).__init__()
              self.lstm = torch.nn.LSTM(input_size, hidden_size, 1)

            def forward(self, x, hidden_state, cell_state):
                # LSTM takes in previous hidden and cell states. The first
                # invokation usually have zero vectors as initial states.
                output, (new_hidden_state, new_cell_state) = \
                    self.lstm(x, (hidden_state, cell_state))
                # LSTM hidden / cell states are returned to be managed by the
                # caller (and is fed in as inputs in the next call).
                return output, new_hidden_state, new_cell_state

        model = TestNet()
        model.eval()

        seq_len = 2 # we'll make seq_len dynamic later
        batch = 1
        input_shape = (seq_len, batch, input_size)
        rand_input = torch.rand(*input_shape)
        h_shape = (1, batch, hidden_size)
        rand_h0 = torch.rand(*h_shape)
        rand_c0 = torch.rand(*h_shape)

        traced_model = torch.jit.trace(model, (rand_input, rand_h0, rand_c0))

        # ct.RangeDim() tells coremltools that this dimension can change for
        # each inference example (aka "runtime-determined"). If the sequence
        # length is always the same (e.g., 2 step LSTM would have seq_len == 2)
        # Note that fixed-length models usually run slightly faster
        # than variable length models.
        ct_seq_len = ct.RangeDim() if variable_length else seq_len
        seq_input = ct.TensorType(shape=(ct_seq_len, batch, input_size),
            name="seq_input")
        h_input = ct.TensorType(shape=h_shape, name="h_input")
        c_input = ct.TensorType(shape=h_shape, name="c_input")

        mlmodel = ct.convert(
            traced_model,
            inputs=[seq_input, h_input, c_input],
        )

        if ct.utils._is_macos():
            result = mlmodel.predict(
                {"seq_input": rand_input.detach().numpy().astype(np.float32),
                  "h_input": rand_h0.detach().numpy().astype(np.float32),
                  "c_input": rand_c0.detach().numpy().astype(np.float32),
                  }
            )

            # Verify outputs
            expected = model(rand_input, rand_h0, rand_c0)
            names = list(result.keys())
            names.sort()
            np.testing.assert_allclose(result[names[0]],
                expected[0].detach().numpy(), atol=1e-4)
            np.testing.assert_allclose(result[names[1]],
                expected[1].detach().numpy(), atol=1e-4)
            np.testing.assert_allclose(result[names[2]],
                expected[2].detach().numpy(), atol=1e-4)

            # Try example of different length
            if variable_length:
                seq_len = 10
                input_shape = (seq_len, batch, input_size)
                rand_input = torch.rand(*input_shape)

                result = mlmodel.predict(
                    {"seq_input": rand_input.detach().numpy().astype(np.float32),
                      "h_input": rand_h0.detach().numpy().astype(np.float32),
                      "c_input": rand_c0.detach().numpy().astype(np.float32),
                      }
                )
                expected = model(rand_input, rand_h0, rand_c0)
                names = list(result.keys())
                names.sort()
                np.testing.assert_allclose(result[names[0]],
                    expected[0].detach().numpy(), atol=1e-4)
                np.testing.assert_allclose(result[names[1]],
                    expected[1].detach().numpy(), atol=1e-4)
                np.testing.assert_allclose(result[names[2]],
                    expected[2].detach().numpy(), atol=1e-4)

    @staticmethod
    @pytest.mark.parametrize(
            "use_symbol", [True, False])
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_outofbound_range_dim(use_symbol):

        num_tokens = 3
        embedding_size = 5

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.embedding = torch.nn.Embedding(num_tokens, embedding_size)

            def forward(self, x):
                return self.embedding(x)

        model = TestModule()
        model.eval()

        example_input = torch.randint(high=num_tokens, size=(3,),
                dtype=torch.int64)
        traced_model = torch.jit.trace(model, example_input)

        if use_symbol:
            seq_len_dim = ct.RangeDim(symbol='len', lower_bound=3,
                    upper_bound=5)
        else:
            # symbol is optional
            seq_len_dim = ct.RangeDim(lower_bound=3, upper_bound=5)
        seq_input = ct.TensorType(name="input", shape=(seq_len_dim,),
                dtype=np.int64)
        mlmodel = ct.convert(
            traced_model,
            inputs=[seq_input],
        )

        if ct.utils._is_macos():
            result = mlmodel.predict(
                {"input": example_input.detach().numpy().astype(np.float32)}
            )

            # Verify outputs
            expected = model(example_input)
            name = list(result.keys())[0]
            np.testing.assert_allclose(result[name], expected.detach().numpy())

            # seq_len below/above lower_bound/upper_bound
            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                example_input2 = torch.randint(high=num_tokens, size=(99,),
                        dtype=torch.int64)
                result = mlmodel.predict(
                    {"input": example_input2.detach().numpy().astype(np.float32)}
                )

            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                example_input2 = torch.randint(high=num_tokens, size=(2,),
                        dtype=torch.int64)
                result = mlmodel.predict(
                    {"input": example_input2.detach().numpy().astype(np.float32)}
                )

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2keras_enumerated_shapes():
        # Test examples in https://coremltools.readme.io/docs/flexible-inputs
        import tensorflow as tf

        input_shape = (28, 28, 3)
        # None denotes seq_len dimension
        x = tf.keras.Input(shape=input_shape, name="input")
        C_out = 2
        kHkW = 3
        y = tf.keras.layers.Conv2D(C_out, kHkW, activation='relu',
                input_shape=input_shape)(x)
        keras_model = tf.keras.Model(inputs=[x], outputs=[y])

        # One RangeDim shared by two inputs
        shapes = [(1, 28, 28, 3), (1, 56, 56, 3)]
        enumerated_shapes = ct.EnumeratedShapes(shapes=shapes)
        tensor_input = ct.TensorType(name="input", shape=enumerated_shapes)
        mlmodel = ct.convert(keras_model, inputs=[tensor_input])

        # Test (1, 28, 28, 3) shape
        test_input_x = np.random.rand(*shapes[0]).astype(np.float32)
        expected_val = keras_model([test_input_x])
        if ct.utils._is_macos():
            results = mlmodel.predict({
                "input": test_input_x})
            np.testing.assert_allclose(results["Identity"],
                    expected_val, rtol=1e-2)

            # Test (1, 56, 56, 3) shape (can't verify numerical parity with Keras
            # which doesn't support enumerated shape)
            test_input_x = np.random.rand(*shapes[1]).astype(np.float32)
            results = mlmodel.predict({
                "input": test_input_x})

            # Test with a wrong shape
            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                test_input_x = np.random.rand(1, 29, 29, 3).astype(np.float32)
                results = mlmodel.predict({
                    "input": test_input_x})

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_enumerated_shapes():

        in_channels = 3
        out_channels = 2
        kernel_size = 3

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.conv = torch.nn.Conv2d(in_channels, out_channels,
                        kernel_size)

            def forward(self, x):
                return self.conv(x)

        model = TestModule()
        model.eval()

        example_input = torch.randn(1, 3, 28, 28)
        traced_model = torch.jit.trace(model, example_input)

        shapes = [(1, 3, 28, 28), (1, 3, 56, 56)]
        enumerated_shapes = ct.EnumeratedShapes(shapes=shapes)
        tensor_input = ct.TensorType(name="input", shape=enumerated_shapes)

        mlmodel = ct.convert(
            traced_model,
            inputs=[tensor_input],
            compute_units=ct.ComputeUnit.CPU_ONLY
        )

        if ct.utils._is_macos():
            result = mlmodel.predict(
                {"input": example_input.detach().numpy().astype(np.float32)},
            )

            # Verify outputs
            expected = model(example_input)
            name = list(result.keys())[0]
            np.testing.assert_allclose(result[name], expected.detach().numpy(),
                    rtol=1e-3, atol=1e-4)

            # Test (1, 3, 56, 56) shape (can't verify numerical parity with Torch
            # which doesn't support enumerated shape)
            test_input_x = np.random.rand(*shapes[1]).astype(np.float32)
            mlmodel.predict({"input": test_input_x})

            # Test with a wrong shape
            with pytest.raises(RuntimeError,
                    match=r"not compatible with the model\'s feature"):
                test_input_x = np.random.rand(1, 3, 29, 29).astype(np.float32)
                mlmodel.predict({"input": test_input_x})

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2_image_enumerated_shapes():
        import tensorflow as tf
        keras_model = tf.keras.applications.MobileNetV2(
            input_shape=(None, None, 3,),
            classes=1000,
            include_top=False,
        )
        input_shapes = ct.EnumeratedShapes(shapes=[(1, 192, 192, 3), (1, 224, 224, 3)])
        image_input = ct.ImageType(shape=input_shapes,
                                   bias=[-1,-1,-1], scale=1/127)
        model = ct.convert(keras_model, inputs=[image_input])
        assert model is not None
        spec = model.get_spec()
        assert len(spec.description.input[0].type.imageType.enumeratedSizes.sizes) == 2

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_image_enumerated_shapes():
        import torchvision
        torch_model = torchvision.models.mobilenet_v2().features
        torch_model.eval()
        example_input = torch.rand(1, 3, 256, 256)
        traced_model = torch.jit.trace(torch_model, example_input)
        input_shapes = ct.EnumeratedShapes(shapes=[(1, 3, 256, 256), (1, 3, 224, 224)])
        image_input = ct.ImageType(shape=input_shapes,
                                   bias=[-1, -1, -1], scale=1 / 127)
        model = ct.convert(traced_model, inputs=[image_input])
        assert model is not None
        spec = model.get_spec()
        assert len(spec.description.input[0].type.imageType.enumeratedSizes.sizes) == 2

    @staticmethod
    @pytest.mark.skipif(ct.utils._macos_version() < (13, 0), reason='Bug fixed in macOS13/iOS16')
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_mean_op():
        # test for bug reported in https://github.com/apple/coremltools/issues/1420
        class Network(torch.nn.Module):
            def forward(self, x):
                return torch.mean(x, dim=(2, 3), keepdim=True)

        model = Network()
        x = torch.rand(1, 3, 256, 256)
        traced_model = torch.jit.trace(model, x)
        input_x = ct.TensorType(shape=(1, 3, ct.RangeDim(default=256), ct.RangeDim(default=256)),
                                name="input")
        cml = ct.convert(traced_model,
                         inputs=[input_x],
                         outputs=[ct.TensorType(name="out")],
                         convert_to="mlprogram",
                         compute_units=ct.ComputeUnit.CPU_ONLY)
        input_dict = {"input": np.random.rand(1, 3, 112, 112)}
        if ct.utils._is_macos():
            out = cml.predict(input_dict)["out"]
            assert out.shape == (1, 3, 1, 1)


class TestOptionalInput:
    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_tf2keras_optional_input():
        # Test examples in https://coremltools.readme.io/docs/flexible-inputs
        import tensorflow as tf

        input_dim = 3
        # None denotes seq_len dimension
        x1 = tf.keras.Input(shape=(None,input_dim), name="optional_input")
        x2 = tf.keras.Input(shape=(None,input_dim), name="required_input")
        y = x1 + x2
        keras_model = tf.keras.Model(inputs=[x1, x2], outputs=[y])

        seq_len_dim = ct.RangeDim()
        default_value = np.ones((1, 2, input_dim)).astype(np.float32)
        optional_input = ct.TensorType(
            name="optional_input",
            shape=(1, seq_len_dim, input_dim),
            default_value=default_value,
          )
        required_input = ct.TensorType(
            name="required_input",
            shape=(1, seq_len_dim, input_dim),
          )
        mlmodel = ct.convert(keras_model,
                inputs=[optional_input, required_input])

        batch = 1
        seq_len = 2
        test_input_x2 = np.random.rand(batch, seq_len, input_dim).astype(np.float32)
        expected_val = keras_model([default_value, test_input_x2])
        if ct.utils._is_macos():
            results = mlmodel.predict({"required_input": test_input_x2})
            np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_classifier():
        class Net(torch.nn.Module):
            def __init__(self):
                super(Net, self).__init__()
                self.linear1 = torch.nn.Linear(28 * 28, 100)
                self.linear2 = torch.nn.Linear(100, 50)
                self.final = torch.nn.Linear(50, 10)
                self.relu = torch.nn.ReLU()

            def forward(self, img):  # convert + flatten
                x = img.view(-1, 28 * 28)
                x = self.relu(self.linear1(x))
                x = self.relu(self.linear2(x))
                x = self.final(x)
                return x
        model = Net()
        model.eval()
        example_input = torch.rand(1, 28 * 28, 1)
        traced_model = torch.jit.trace(model, example_input)
        traced_model.eval()

        def _test_classifier(traced_model, example_input, class_type, backend):
            label = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
            if class_type == "str":
                label = list(map(lambda x: str(x), label))
            classifier_config = ct.ClassifierConfig(label)
            mlmodel = ct.convert(
                traced_model,
                source='pytorch',
                convert_to=backend,
                inputs=[
                    ct.TensorType(
                        name="input",
                        shape=example_input.shape,
                        dtype=example_input.numpy().dtype,
                    )
                ],
                classifier_config=classifier_config
            )
            if ct.utils._is_macos():
                coreml_out = mlmodel.predict({"input": example_input.detach().numpy()})
                assert "classLabel" in coreml_out
                key_type = str if class_type == "str" else int
                assert isinstance(coreml_out["classLabel"], key_type)

        for class_type in ("str", "int"):
            _test_classifier(traced_model, example_input, class_type, "neuralnetwork")
            if ct.utils._macos_version() >= (12, 0):
                _test_classifier(traced_model, example_input, class_type, "mlprogram")

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_torch_optional_input():

        num_tokens = 3
        embedding_size = 5

        class TestModule(torch.nn.Module):
            def __init__(self):
                super(TestModule, self).__init__()
                self.embedding = torch.nn.Embedding(num_tokens, embedding_size)

            def forward(self, x, y):
                return self.embedding(x) + y

        model = TestModule()
        model.eval()

        example_input = [
            torch.randint(high=num_tokens, size=(2,), dtype=torch.int64),
            torch.rand(1),
            ]
        traced_model = torch.jit.trace(model, example_input)

        required_input = ct.TensorType(
            name="required_input", shape=(ct.RangeDim(),), dtype=np.int64)
        default_value = np.array([3]).astype(np.float32)
        optional_input = ct.TensorType(name="optional_input", shape=(1,),
            default_value=default_value)

        for compute_units in ct.ComputeUnit:
            if compute_units == ct.ComputeUnit.CPU_AND_NE and ct.utils._macos_version() < (13, 0):
                continue

            mlmodel = ct.convert(
                traced_model,
                inputs=[required_input, optional_input],
                compute_units=compute_units,
            )

            assert(mlmodel.compute_unit == compute_units)

            if ct.utils._is_macos():
                result = mlmodel.predict(
                    {"required_input":
                     example_input[0].detach().numpy().astype(np.float32)}
                )

                # Verify outputs
                torch_default_value = torch.tensor([3])
                expected = model(example_input[0].detach(), torch_default_value)
                name = list(result.keys())[0]
                np.testing.assert_allclose(result[name], expected.detach().numpy())

###############################################################################
# Note: all tests are examples of conversion to the Core ML format
# Each test case is expected to be runnable and self-complete.
###############################################################################

@pytest.mark.skipif(ct.utils._macos_version() < (10, 15), reason='Model produces specification 4.')
class TestMILConverterExamples:

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND)
    def test_convert_tf1_frozen_graph(tmpdir):
        import tensorflow as tf

        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input")
            y = tf.nn.relu(x, name="output")

        model = ct.convert(graph, convert_to='milinternal')
        assert isinstance(model, ct.converters.mil.Program)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND)
    def test_convert_tf2_keras(tmpdir):
        import tensorflow as tf

        x = tf.keras.Input(shape=(32,), name="input")
        y = tf.keras.layers.Dense(16, activation="softmax")(x)
        keras_model = tf.keras.Model(x, y)
        model = ct.convert(keras_model, convert_to='milinternal')
        assert isinstance(model, ct.converters.mil.Program)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_convert_torch_traced_model(tmpdir):
        from torch import nn
        class Network(nn.Module):
            def __init__(self):
                super(Network, self).__init__()
                self.hidden = nn.Linear(100, 10)
                self.output = nn.Linear(10, 2)
                self.sigmoid = nn.Sigmoid()
                self.softmax = nn.Softmax(dim=1)

            def forward(self, x):
                x = self.hidden(x)
                x = self.sigmoid(x)
                x = self.output(x)
                x = self.softmax(x)
                return x

        torch_model = Network()
        torch_model.eval()
        example_input = torch.rand(1, 100)
        traced_model = torch.jit.trace(torch_model, example_input)
        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
            convert_to='milinternal'
        )
        assert isinstance(model, ct.converters.mil.Program)


    @staticmethod
    @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND)
    def test_mil_op_names_consistency(tmpdir):
        '''
        Test to make sure that when the same model is converted to MIL program,
        in the same session, it gives the same program, with the same op names
        '''
        import tensorflow as tf

        with tf.Graph().as_default() as graph:
            x = tf.placeholder(tf.float32, shape=(1, 5, 5, 3), name="input")
            conv = tf.nn.conv2d(
                x,
                filter = tf.constant(np.random.rand(1, 1, 3, 5), tf.float32),
                padding = "VALID",
            )
            y = tf.nn.relu(conv, name="output")

        mil_prog1 = ct.convert(graph, convert_to='milinternal')
        # convert the same model again
        mil_prog2 = ct.convert(graph, convert_to='milinternal')

        # compare op names of the two programs
        np.testing.assert_array_equal(get_op_types_in_program(mil_prog1), get_op_types_in_program(mil_prog2))

class TestMLProgramConverterExamples:
    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    @pytest.mark.parametrize(
        "convert_to", ['neuralnetwork', 'mlprogram'])
    def test_convert_to_argument(tmpdir, convert_to):
        class Network(torch.nn.Module):
            def __init__(self):
                super(Network, self).__init__()
                self.hidden = torch.nn.Linear(30, 5)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                x = self.hidden(x)
                return self.relu(x)

        torch_model = Network()
        torch_model.eval()
        example_input = torch.rand(1, 30)
        traced_model = torch.jit.trace(torch_model, example_input)
        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
            convert_to=convert_to
        )
        assert isinstance(model, ct.models.MLModel)
        spec = model.get_spec()
        if convert_to == "mlprogram":
            assert spec.WhichOneof('Type') == 'mlProgram'
        else:
            assert spec.WhichOneof('Type') == 'neuralNetwork'

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_deployment_target_argument(tmpdir):
        class Network(torch.nn.Module):
            def __init__(self):
                super(Network, self).__init__()
                self.hidden = torch.nn.Linear(30, 5)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                x = self.hidden(x)
                return self.relu(x)

        torch_model = Network()
        torch_model.eval()
        example_input = torch.rand(1, 30)
        traced_model = torch.jit.trace(torch_model, example_input)

        # convert to 'neuralnetwork' by specifying an iOS13 target
        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
            minimum_deployment_target=ct.target.iOS13,
        )
        assert isinstance(model, ct.models.MLModel)
        assert model.get_spec().WhichOneof('Type') == 'neuralNetwork'

        # convert to 'mlprogram' by specifying an iOS15 target
        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
            minimum_deployment_target=ct.target.iOS15,
        )
        assert isinstance(model, ct.models.MLModel)
        assert model.get_spec().WhichOneof('Type') == 'mlProgram'

        # verify an error is raised when convert_to="neuralnetwork" and target is iOS15
        with pytest.raises(ValueError) as e:
            model = ct.convert(
                traced_model,
                inputs=[ct.TensorType(name="input", shape=example_input.shape)],
                convert_to="neuralnetwork",
                minimum_deployment_target=ct.target.iOS15,
            )
        expected_error = "If minimum deployment target is iOS15/macOS12/watchOS8/tvOS15 or higher, " \
                         "then 'convert_to' cannot be neuralnetwork. It must be 'mlprogram'"
        assert expected_error == str(e.value)

        # verify an error is raised when convert_to="mlprogram" and target is less than iOS15
        with pytest.raises(ValueError) as e:
            model = ct.convert(
                traced_model,
                inputs=[ct.TensorType(name="input", shape=example_input.shape)],
                convert_to="mlprogram",
                minimum_deployment_target=ct.target.iOS14,
            )
        expected_error = "When 'convert_to' is mlprogram, the minimum deployment target " \
                         "must be at least iOS15/macOS12/watchOS8/tvOS15"
        assert expected_error == str(e.value)

    @staticmethod
    def test_model_save(tmpdir):
        save_path_dir = str(tmpdir)

        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
        def prog(x):
            x = mb.square(x=x)
            return x

        # save neuralnetwork model without extension and check that it is saved with
        # mlmodel extension
        mlmodel = ct.convert(prog)
        mlmodel_path = os.path.join(save_path_dir, "model_nn")
        mlmodel.save(mlmodel_path)
        assert os.path.exists(mlmodel_path + ".mlmodel")

        # save neuralnetwork model with mlpackage extension
        mlmodel_path = os.path.join(save_path_dir, "model_nn2.mlpackage")
        mlmodel.save(mlmodel_path)
        assert os.path.exists(mlmodel_path)

        # save mlprogram model without extension and check that it is saved with
        # mlpackage extension
        mlmodel = ct.convert(prog, convert_to="mlprogram")
        mlmodel_path = os.path.join(save_path_dir, "model_mlprogram")
        mlmodel.save(mlmodel_path)
        assert os.path.exists(mlmodel_path + ".mlpackage")

        # check error if mlprogram is saved with mlmodel extension
        mlmodel_path = os.path.join(save_path_dir, "model_mlprogram.mlmodel")
        with pytest.raises(Exception) as e:
            mlmodel.save(mlmodel_path)
        expected_error = "For an ML Program, extension must be .mlpackage (not .mlmodel)"
        assert expected_error == str(e.value)

    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
    def test_get_milprogram_method(tmpdir):
        class Network(torch.nn.Module):
            def __init__(self):
                super(Network, self).__init__()
                self.hidden = torch.nn.Linear(100, 10)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                x = self.hidden(x)
                x = self.relu(x)
                return x

        torch_model = Network()
        torch_model.eval()
        example_input = torch.rand(1, 100)
        traced_model = torch.jit.trace(torch_model, example_input)
        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(shape=example_input.shape)],
            convert_to='mlprogram'
        )
        assert isinstance(model._get_mil_internal(), ct.converters.mil.Program)

    @staticmethod
    @pytest.mark.skipif(not ct.utils._is_macos(), reason="Platform is not Mac OS")
    def test_deepcopy_error_with_symbols_in_prog():
        prog = Program()
        func_inputs = {"x": mb.placeholder(shape=[get_new_symbol(), 3]),
                       "y": mb.placeholder(shape=[2, 3])}
        with Function(func_inputs) as ssa_fun:
            x, y = ssa_fun.inputs["x"], ssa_fun.inputs["y"]
            x = mb.relu(x=x)
            z = mb.add(x=x, y=y)
            ssa_fun.set_outputs([z])
        prog.add_function("main", ssa_fun)
        mlmodel = ct.convert(prog, convert_to="mlprogram", compute_precision=ct.precision.FLOAT32)
        prog2 = mlmodel._get_mil_internal() # this will invoke a deepcopy on the prog


    @staticmethod
    @pytest.mark.skipif(not _HAS_TORCH or ct.utils._macos_version() < (12, 0),
                        reason=MSG_TORCH_NOT_FOUND)
    @pytest.mark.parametrize(
        "convert_to, provide_prob_output_argument",
        itertools.product(
            ["neuralnetwork", "mlprogram"],
            [False, True],
        )
    )
    def test_classifier(convert_to, provide_prob_output_argument):
        torch_model = torch.nn.ReLU().eval()
        traced_model = torch.jit.trace(torch_model, torch.rand(3,))
        variable_name = "var_2"
        class_label_name = "class_label"
        classifier_config = ct.ClassifierConfig(
                            class_labels=['a', 'b', 'c'],
                            predicted_feature_name=class_label_name,
                            predicted_probabilities_output=variable_name if provide_prob_output_argument else None,
                            )

        model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(shape=(3,))],
            classifier_config = classifier_config,
            convert_to=convert_to,
        )
        spec = model.get_spec()
        input_name = spec.description.input[0].name
        out_dict = model.predict({input_name : np.array([1.0, 2.0, 3.0])})

        assert class_label_name in out_dict
        assert out_dict[class_label_name] == 'c'
        if convert_to == "neuralnetwork":
            assert variable_name in out_dict
            assert isinstance(out_dict[variable_name], dict)
        else:
            output_dict_feature_name = class_label_name + "_probs"
            assert output_dict_feature_name in out_dict
            assert isinstance(out_dict[output_dict_feature_name], dict)

    @pytest.mark.skipif(not ct.utils._is_macos(), reason="Platform is not Mac OS")
    @pytest.mark.parametrize("skip_model_load", [True, False])
    def test_model_load_skip_flag(self, skip_model_load):
        @mb.program(input_specs=[mb.TensorSpec(shape=(3,)), ])
        def prog(x):
            return mb.relu(x=x, name='relu')

        if ct.utils._macos_version() < (12, 0) and not skip_model_load:
            # converting to mlprogram, on macOS < 12
            # should raise a runtime error when skip_model_load is False
            with pytest.warns(RuntimeWarning):
                model = ct.convert(prog, convert_to='mlprogram',
                                   skip_model_load=skip_model_load)
        else:
            model = ct.convert(prog, convert_to='mlprogram',
                                skip_model_load=skip_model_load)

        assert model is not None
        if skip_model_load:
            assert model.__proxy__ is None
        model_dir = mkdtemp()
        filename = os.path.join(model_dir, 'test.mlpackage')
        model.save(filename)
        assert os.path.exists(filename)
        try:
            rmtree(model_dir)
        except:
            pass


@pytest.mark.skipif(ct.utils._macos_version() < (12, 0), reason='Model produces specification 6.')
class TestMLProgramFP16Transform:
    @staticmethod
    def test_compute_precision_api():
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
        def prog(x):
            x = mb.square(x=x)
            return x

        mlmodel = ct.convert(copy.deepcopy(prog),
                             compute_precision=ct.precision.FLOAT16,
                             convert_to='mlprogram')
        mil_prog = mlmodel._get_mil_internal()
        np.testing.assert_array_equal(["cast", "square", "cast"], get_op_types_in_program(mil_prog))

        mlmodel = ct.convert(copy.deepcopy(prog),
                             compute_precision=ct.precision.FLOAT32,
                             convert_to='mlprogram')
        mil_prog = mlmodel._get_mil_internal()
        np.testing.assert_array_equal(["square"], get_op_types_in_program(mil_prog))

        mlmodel = ct.convert(copy.deepcopy(prog),
                             compute_precision=ct.transform.FP16ComputePrecision(
                                                            op_selector=lambda op: op.op_type != 'square'),
                             convert_to='mlprogram')
        mil_prog = mlmodel._get_mil_internal()
        np.testing.assert_array_equal(["square"], get_op_types_in_program(mil_prog))

        with pytest.raises(ValueError) as e:
            mlmodel = ct.convert(copy.deepcopy(prog),
                                 compute_precision='fp64',
                                 convert_to='mlprogram')
        expected_error = "'compute_precision' must be either coremltools.precision.FLOAT32 or " \
                         "coremltools.precision.FLOAT16 or of type coremltools.transform.FP16ComputePrecision()"
        assert expected_error == str(e.value)

        expected_pattern = "compute_precision .* supported .* mlprogram .* None .* target==\'neuralnetwork\'.*\n.*minimum_deployment_target.*"
        with pytest.raises(ValueError, match=expected_pattern) as e:
            mlmodel = ct.convert(copy.deepcopy(prog), compute_precision='fp16')

    @staticmethod
    def test_invalid_argument_nn_backend():
        '''
        Since the  compute_precision argument is only applicable when converting to "mlprogram",
        check that an error is correctly raised when conversion is targeted at the neuralnetwork backend
        '''
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
        def prog(x):
            x = mb.square(x=x)
            return x

        expected_err_str = "compute_precision is only supported for mlprogram target and must be None if target.*"
        with pytest.raises(ValueError, match=expected_err_str):
            mlmodel = ct.convert(prog, compute_precision=ct.precision.FLOAT16)
        with pytest.raises(ValueError, match=expected_err_str):
            mlmodel = ct.convert(prog, compute_precision=ct.precision.FLOAT32)
