code: def test_layernorm_fusion(axes_size):
      """
      Detect layer norm pattern, found in the TF bert model.
      y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)])
  
      where mean and variance are computed along axes [-1] or [-1,-2] and so on
      and gamma and beta are constants with rank equal to the length of the axes parameter.
      """
      shape = (3, 5, 6)
      rank = len(shape)
      axes = list(range(rank - axes_size, rank))
  
      @mb.program(input_specs=[mb.TensorSpec(shape=shape)])
      def prog(x):
          x1 = mb.reduce_mean(x=x, axes=axes, keep_dims=True)
          x2 = mb.sub(x=x, y=x1)
          x2 = mb.square(x=x2)
          x2 = mb.reduce_mean(x=x2, axes=axes, keep_dims=True)
          x2 = mb.add(x=x2, y=1e-5)
          x2 = mb.rsqrt(x=x2)
          x3 = mb.mul(x=np.random.rand(*shape[-len(axes) :]), y=x2)
          x4 = mb.mul(x=x3, y=x1)
          x5 = mb.mul(x=x, y=x3)
          x4 = mb.sub(x=np.random.rand(*shape[-len(axes) :]), y=x4)
          y = mb.add(x=x4, y=x5)
          return y
  
      prev_prog, prev_block, block = apply_pass_and_basic_check(
          prog, "common::fuse_layernorm_or_instancenorm"
      )
      assert get_op_types_in_program(prev_prog) == [
          "reduce_mean",
          "sub",
          "square",
          "reduce_mean",
          "add",
          "rsqrt",
          "mul",
          "mul",
          "mul",
          "sub",
          "add",
      ]
      assert get_op_types_in_program(prog) == ["layer_norm"]
      assert_model_is_valid(
          prog,
          {"x": shape},
          expected_output_shapes={block.outputs[0].name: shape},
      )
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/mil/passes/test_passes.py
imports: from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.testing_utils import assert_op_count_match, assert_model_is_valid, assert_same_output_names, get_op_types_in_program, apply_pass_and_basic_check
  from coremltools.converters.mil.mil import Symbol, types
  from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
  import copy
  import pytest
  import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: test_layernorm_fusion(axes_size) and returns: y