text: #  Copyright (c) 2020, 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 pytest
  import itertools
  import numpy as np
  import copy
  
  # import mil internal ops to add it to the builder
  import coremltools as ct
  from coremltools.converters.mil.mil import Builder as mb
  from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
  from coremltools.converters.mil import types
  from coremltools.converters.mil.mil.types import string_to_builtin, builtin_to_string, promote_types
  
  # Set the testing backend
  import coremltools.converters.mil.testing_reqs as testing_reqs
  
  from coremltools.converters.mil.testing_utils import (
      get_op_types_in_program,
      apply_pass_and_basic_check,
      assert_model_is_valid,
  )
  
  
  class TestAdjustToSupportedTypes:
  
      def test_basic(self):
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.bool),
                                   mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int32),
                                   mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.fp32)])
          def prog(x, y, z):
              out = mb.logical_not(x=x)
              return (out, y, z)
          prog.functions['not_main'] = copy.deepcopy(prog.functions['main'])
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          """
          Input graph:
  
          func main(bool x, int32 y, fp32 z) {
              bool  out = logical_not(x)
          } -> (out, y, z)
  
          becomes
  
          func main(fp32 x, int32 y, fp32 z) {
             bool  x_casted = cast(x)
             bool  out__pre__output__fp32__cast = logical_not(x_casted)
             fp32  out = cast(out__pre__output__fp32__cast)
          } -> (out, y, z)
          """
          assert get_op_types_in_program(prev_prog) == ['logical_not']
          assert get_op_types_in_program(prog) == ['cast', 'logical_not', 'cast']
  
          prev_inputs = list(prev_prog.functions['main'].inputs.items())
          inputs = list(prog.functions['main'].inputs.items())
          assert prev_inputs[0][1].name == inputs[0][1].name
          assert inputs[0][1].dtype == types.fp32
          for i in range(1, len(inputs)):
              assert prev_inputs[i][1].name == inputs[i][1].name
              assert prev_inputs[i][1].dtype == inputs[i][1].dtype
  
          prev_outputs = prev_prog.functions['main'].outputs
          outputs = prog.functions['main'].outputs
          assert prev_outputs[0].name == outputs[0].name
          assert outputs[0].dtype == types.fp32
          for i in range(1, len(outputs)):
              assert prev_outputs[i].name == outputs[i].name
              assert prev_outputs[i].dtype == outputs[i].dtype
  
          """
          Input graph:
  
          func not_main(bool x, int32 y, fp32 z) {
              bool  out = logical_not(x)
          } -> (out, y, z)
  
          is identical after the pass.
          """
          assert get_op_types_in_program(prev_prog, 'not_main') == ['logical_not']
          assert get_op_types_in_program(prog, 'not_main') == ['logical_not']
  
          prev_inputs = list(prev_prog.functions['not_main'].inputs.items())
          inputs = list(prog.functions['not_main'].inputs.items())
          for i in range(0, len(inputs)):
              assert prev_inputs[i][1].name == inputs[i][1].name
              assert prev_inputs[i][1].dtype == inputs[i][1].dtype
  
          prev_outputs = prev_prog.functions['not_main'].outputs
          outputs = prog.functions['not_main'].outputs
          for i in range(0, len(outputs)):
              assert prev_outputs[i].name == outputs[i].name
              assert prev_outputs[i].dtype == outputs[i].dtype
  
      def test_int64_input(self):
          """
          Input graph:
  
          func main(int64 x) {
          } -> (x)
  
          becomes
  
          func main(int32 x) {
          } -> (x)
          """
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int64)])
          def prog(x):
              return x
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          prev_inputs = list(prev_prog.functions['main'].inputs.items())
          inputs = list(prog.functions['main'].inputs.items())
          assert prev_inputs[0][1].name == inputs[0][1].name
          assert inputs[0][1].dtype == types.int32
  
      def test_float64_input(self):
          """
          Input graph:
  
          func main(float64 x) {
          } -> (x)
  
          becomes
  
          func main(float32 x) {
          } -> (x)
          """
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.fp64)])
          def prog(x):
              return x
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          prev_inputs = list(prev_prog.functions['main'].inputs.items())
          inputs = list(prog.functions['main'].inputs.items())
          assert prev_inputs[0][1].name == inputs[0][1].name
          assert inputs[0][1].dtype == types.fp32
  
      def test_int8_input(self):
          """
          Input graph:
  
          func main(int8 x) {
          } -> (x)
  
          becomes
  
          func main(int32 x) {
          } -> (x)
          """
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int8)])
          def prog(x):
              return x
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          prev_inputs = list(prev_prog.functions['main'].inputs.items())
          inputs = list(prog.functions['main'].inputs.items())
          assert prev_inputs[0][1].name == inputs[0][1].name
          assert inputs[0][1].dtype == types.int32
  
      def test_subblock(self):
          """
          Input graph:
  
          func main(float64 a, float32 b) {
              float64 out_0, float32 out_1 = while_loop(a, b,
                  (float64 a, float32 b) {
                      bool cond = less(a, b)
                  } -> (cond)
                  (float64 a, float32 b) {
                      float64 temp = const(1)
                      float64 out = add(a, b)
                  } -> (out, b)
              );
          } -> (out_0, out_1)
  
          becomes
  
          func main(float32 a, float32 b) {
              float32 out_0, float32 out_1 = while_loop(a, b,
                  (float32 a, float32 b) {
                      bool cond = less(a, b)
                  } -> (cond)
                  (float32 a, float32 b) {
                      float32 temp = const(1)
                      float32 out = add(a, b)
                  } -> (out, b)
              );
          } -> (out_0, out_1)
          """
          def body(a, b):
              return mb.add(x=a, y=np.float64(1)), b
  
          def cond(a, b):
              return mb.less(x=a, y=b)
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=types.fp64),
                                   mb.TensorSpec(shape=(1,), dtype=types.fp32)])
          def prog(a, b):
              return mb.while_loop(_cond=cond, _body=body, loop_vars=(a, b))
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          prev_inputs = list(prev_prog.functions['main'].inputs.items())
          inputs = list(prog.functions['main'].inputs.items())
          for i in range(0, len(prev_inputs)):
              assert prev_inputs[i][1].name == inputs[i][1].name
              assert inputs[i][1].dtype == types.fp32
  
          assert get_op_types_in_program(prev_prog) == ['while_loop']
          assert get_op_types_in_program(prog) == ['while_loop']
  
          def assert_block_inputs(prev_inputs, inputs):
              for i in range(0, len(prev_inputs)):
                  assert prev_inputs[i].name == inputs[i].name
                  assert inputs[i].dtype == types.fp32
  
          subblocks = prog.functions['main'].operations[0].blocks
          prev_subblocks = prev_prog.functions['main'].operations[0].blocks
          for i in range(0, len(subblocks)):
              assert_block_inputs(prev_subblocks[i].inputs, subblocks[i].inputs)
  
      def test_adjust_cast(self):
          """
          Input graph:
  
          func main(int32 x) {
              fp64 y = cast(x=x, dtype="fp64")
          } -> (y)
  
          becomes
  
          func main(int32 x) {
              fp32 y = cast(x=x, dtype="fp32")
          } -> (y)
          """
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int32)])
          def prog(x):
              y = mb.cast(x=x, dtype="fp64")
              return y
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          assert get_op_types_in_program(prev_prog) == ['cast']
          assert get_op_types_in_program(prog) == ['cast']
  
          prev_cast = prev_prog.functions['main'].operations[1]
          cast = prog.functions['main'].operations[2]
  
          assert prev_cast.dtype.val == "fp64"
          assert prev_cast.outputs[0].dtype == types.fp64
  
          assert cast.dtype.val == "fp32"
          assert cast.outputs[0].dtype == types.fp32
  
      def test_adjust_redundant_cast(self):
          """
          Input graph:
  
          func main(int32 x) {
              int64 y = cast(x=x, dtype="int64")
          } -> (y)
  
          becomes
  
          func main(int32 x) {
          } -> (x)
          """
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int32)])
          def prog(x):
              y = mb.cast(x=x, dtype="int64")
              return y
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::adjust_io_to_supported_types"
          )
  
          assert get_op_types_in_program(prev_prog) == ['cast']
          assert get_op_types_in_program(prog) == []
  
  class TestImagePreprocessingPass:
  
      def test_program_grayscale(self):
          """
          Input graph:
  
          main(x: ImageType(color_layout="G", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 1, 20, 20],
                                                color_layout="G",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["relu", "relu", "add"]
  
      def test_program_grayscale_with_scale(self):
          """
          Input graph:
  
          main(x: ImageType(scale=2.0, color_layout="G", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y = mul(x, 2)
              y1 = relu(y)
              y2 = relu(y)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 1, 20, 20],
                                                scale=2.0,
                                                color_layout="G",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["mul", "relu", "relu", "add"]
          scale_op = prog.find_ops(op_type="mul", exactly_one=True)[0]
          assert scale_op.y.val == 2.0
  
      def test_program_grayscale_with_bias(self):
          """
          Input graph:
  
          main(x: ImageType(bias=2.0, color_layout="G", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y = add(x, 2)
              y1 = relu(y)
              y2 = relu(y)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 1, 20, 20],
                                                bias=2.0,
                                                color_layout="G",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["add", "relu", "relu", "add"]
          add_op = prog.find_ops(op_type="add", exactly_one=False)[0]
          assert add_op.y.val == 2.0
  
      def test_program_grayscale_with_scale_bias(self):
          """
          Input graph:
  
          main(x: ImageType(scale=2.0, bias=2.0, color_layout="G", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y_scaled = mul(x, 2)
              y = add(y_scaled, 2)
              y1 = relu(y)
              y2 = relu(y)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 1, 20, 20],
                                                scale=2.0,
                                                bias=2.0,
                                                color_layout="G",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["mul", "add", "relu", "relu", "add"]
          scale_op = prog.find_ops(op_type="mul", exactly_one=True)[0]
          assert scale_op.y.val == 2.0
          add_op = prog.find_ops(op_type="add", exactly_one=False)[0]
          assert add_op.y.val == 2.0
  
      def test_program_rgb(self):
          """
          Input graph:
  
          main(x: ImageType(color_layout="RGB", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 3, 20, 20],
                                                color_layout="RGB",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["relu", "relu", "add"]
  
      def test_program_rgb_scale_bias(self):
          """
          Input graph:
  
          main(x: ImageType(color_layout="RGB", scale=2.0, bias=[1.0, 2.0, 3.0], channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y = mul(x, scale)
              y_bias = add(y, bias)
              y1 = relu(y_bias)
              y2 = relu(y_bias)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 3, 20, 20],
                                                scale=2.0,
                                                bias=[1.0, 2.0, 3.0],
                                                color_layout="RGB",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["mul", "add", "relu", "relu", "add"]
          scale_op = prog.find_ops(op_type="mul", exactly_one=True)[0]
          assert scale_op.y.val == 2.0
          add_op = prog.find_ops(op_type="add", exactly_one=False)[0]
          assert np.all(add_op.y.val == np.array([1.0, 2.0, 3.0]).reshape([1, 3, 1, 1]))
  
      def test_program_bgr(self):
          """
          Input graph:
  
          main(x: ImageType(color_layout="BGR", channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 3, 20, 20],
                                                color_layout="BGR",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["relu", "relu", "add"]
  
      def test_program_bgr_scale_bias(self):
          """
          Input graph:
  
          main(x: ImageType(color_layout="BGR", scale=2.0, bias=[1.0, 2.0, 3.0], channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y = mul(x, scale)
              y_bias = add(y, bias)
              y1 = relu(y_bias)
              y2 = relu(y_bias)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 3, 20, 20],
                                                scale=2.0,
                                                bias=[1.0, 2.0, 3.0],
                                                color_layout="BGR",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["mul", "add", "relu", "relu", "add"]
          scale_op = prog.find_ops(op_type="mul", exactly_one=True)[0]
          assert scale_op.y.val == 2.0
          add_op = prog.find_ops(op_type="add", exactly_one=False)[0]
          assert np.all(add_op.y.val == np.array([1.0, 2.0, 3.0]).reshape([1, 3, 1, 1]))
  
      @pytest.mark.parametrize(
          "scale_type, bias_type", itertools.product([np.float, np.int32], [np.float, np.int32])
      )
      def test_scale_bias_types(self, scale_type, bias_type):
          """
          Input graph:
  
          main(x: ImageType(color_layout="RGB", scale=2.0, bias=[1.0, 2.0, 3.0], channel_first=True)) {
              y1 = relu(x)
              y2 = relu(x)
              output = add(y1, y2)
          } [output]
  
          Output graph:
  
          main(x: ImageType(channel_first=True)) {
              y = mul(x, scale)
              y_bias = add(y, bias)
              y1 = relu(y_bias)
              y2 = relu(y_bias)
              output = add(y1, y2)
          } [output]
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
          def prog(x):
              y1 = mb.relu(x=x)
              y2 = mb.relu(x=x)
              z = mb.add(x=y1, y=y2)
              return z
  
          prog.main_input_types = (ct.ImageType(name='x',
                                                shape=[1, 3, 20, 20],
                                                scale=scale_type(2.0),
                                                bias=np.array([1, 2, 3]).astype(bias_type),
                                                color_layout="RGB",
                                                channel_first=True),)
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              prog, "mil_backend::insert_image_preprocessing_ops"
          )
          assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
          assert get_op_types_in_program(prog) == ["mul", "add", "relu", "relu", "add"]
          scale_op = prog.find_ops(op_type="mul", exactly_one=True)[0]
          assert scale_op.y.dtype() == prog.functions["main"].inputs["x"].dtype()
          add_op = prog.find_ops(op_type="add", exactly_one=False)[0]
          assert add_op.y.dtype() == prog.functions["main"].inputs["x"].dtype()
  
  class TestSanitizerPass:
  
      def test_sanitize_numeric_var_names(self):
          """
          Input:
              main(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1!: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1!")
                  %1: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="1")
                  %3: (1, 3, 20, fp32)(Tensor) = add(x=%Var_1!, y=%1, name="3")
                } -> (%3)
              }
  
          Output:
              main(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1_: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1_")
                  %var_1: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="op_1")
                  %var_3: (1, 3, 20, fp32)(Tensor) = add(x=%var_1_, y=%var_1, name="op_3")
                } -> (%var_3)
              }
  
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20))])
          def prog(x):
              y1 = mb.relu(x=x, name = "var_1!")
              y2 = mb.relu(x=x, name = "1")
              z = mb.add(x=y1, y=y2, name = "3")
              return z
  
          PASS_REGISTRY["mil_backend::sanitize_name_strings"](prog)
          block = prog.functions["main"]
          assert block.find_ops(op_type="relu")[0].outputs[0].name == "var_1_"
          assert block.find_ops(op_type="relu")[1].outputs[0].name == "var_1"
          assert prog["main"].outputs[0].name == "var_3"
          assert block.find_ops(op_type="relu")[0].name == "var_1_"
          assert block.find_ops(op_type="relu")[1].name == "op_1"
          assert block.find_ops(op_type="add")[0].name == "op_3"
  
      def test_sanitize_var_names_with_two_functions(self):
          """
          Input:
              main(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1!: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1!")
                } -> (%var_1!)
              }
  
              main_2(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1!: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1!")
                } -> (%var_1!)
              }
  
  
          Output:
              main(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1!: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1_")
                } -> (%var_1_)
              }
  
              main_2(%x: (1, 3, 20, fp32)(Tensor)) {
                block0() {
                  %var_1!: (1, 3, 20, fp32)(Tensor) = relu(x=%x, name="var_1_")
                } -> (%var_1_)
              }
  
          """
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20))])
          def prog(x):
              z = mb.relu(x=x, name = "var_1!")
              return z
  
          @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20))])
          def prog_2(x):
              z = mb.relu(x=x, name = "var_1!")
              return z
  
          prog.add_function("main_2", prog_2.functions["main"])
          PASS_REGISTRY["mil_backend::sanitize_name_strings"](prog)
          block = prog.functions["main"]
          assert block.find_ops(op_type="relu")[0].outputs[0].name == "var_1_"
          assert prog["main"].outputs[0].name == "var_1_"
          assert block.find_ops(op_type="relu")[0].name == "var_1_"
          block = prog.functions["main_2"]
          assert block.find_ops(op_type="relu")[0].outputs[0].name == "var_1_"
          assert prog["main"].outputs[0].name == "var_1_"
          assert block.find_ops(op_type="relu")[0].name == "var_1_"
  
  
  class TestPassFuseActivationSiLU:
      """
      Input graph:
      input --> sigmoid --> mul --> output
      Output graph:
      input --> silu --> output
      """
  
      @pytest.mark.skipif(ct.utils._macos_version() < (12, 0), reason="mlprogram predict available only on macOS12+")
      @pytest.mark.parametrize(
          "reverse_order", itertools.product([True, False]),
      )
      def test_0(self, reverse_order):
          x_shape = tuple(np.random.randint(low=1, high=4, size=5))
  
          @mb.program(input_specs=[mb.TensorSpec(shape=x_shape)])
          def program(x):
              sigmoid_x = mb.sigmoid(x=x)
              if not reverse_order:
                  x = mb.mul(x=x, y=sigmoid_x)
              else:
                  x = mb.mul(x=sigmoid_x, y=x)
              return x
  
          prev_prog, prev_block, block = apply_pass_and_basic_check(
              program, "mil_backend::fuse_activation_silu"
          )
  
          assert get_op_types_in_program(prev_prog) == ["sigmoid", "mul"]
          assert get_op_types_in_program(program) == ["silu"]
  
          assert_model_is_valid(
              program=program,
              inputs={"x": x_shape},
              backend=("mlprogram", "fp32"),
              expected_output_shapes={block.outputs[0].name: tuple(x_shape)},
          )
  
  
  class TestHomogenizeInputDtypes:
  
      @pytest.mark.parametrize(
          ["op", "x_dtype", "y_dtype"],
          [
              ["add", "int32", "fp32"],
              ["mul", "fp32", "int32"],
              ["minimum", "int64", "int32"],
              ["add", "int32", "fp16"],
              ["add", "fp16", "int32"],
              ["equal", "bool", "int32"],
              ["mod", "int64", "fp16"],
              ["not_equal", "fp32", "bool"],
              ["pow", "fp16", "fp32"],
              ["greater", "fp16", "fp32"],
              ["matmul", "fp16", "int32"],
          ]
      )
      def test_mixed_input_dtypes(self, op, x_dtype, y_dtype):
          @mb.program(input_specs=[mb.TensorSpec(shape=(10, 10), dtype=string_to_builtin(x_dtype)),
                                   mb.TensorSpec(shape=(10, 10), dtype=string_to_builtin(y_dtype))])
          def prog(x, y):
              x = getattr(mb, op)(x=x, y=y)
              return x
  
          assert get_op_types_in_program(prog) == [op]
  
          _, _, block = apply_pass_and_basic_check(prog, "mil_backend::homogenize_input_dtypes")
  
          assert get_op_types_in_program(prog) == ["cast", op]
  
          promoted_dtype = promote_types(string_to_builtin(x_dtype), string_to_builtin(y_dtype))
  
          # Asserting cast configuration
          cast = block.find_ops(op_type="cast")[0]
          assert cast.dtype.val == builtin_to_string(promoted_dtype)
          assert len(cast.outputs) == 1
          assert len(cast.outputs[0].child_ops) == 1
          assert cast.outputs[0].child_ops[0].op_type == op
  