code: def fuse_layernorm_or_instancenorm(prog):
      """
      Identify the pattern:
  
      y = gamma * (x - mean) / sqrt(variance + epsilon) + beta
  
      y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)])
  
      [....] ---> reduce_mean --->  sub -->square -->reduce_mean --> add(epsilon)-->rsqrt
         |             |             ^                                                |
         |             |             |                                                |
         |             |             |                                                V
         |---------------------------                                               mul (gamma)
         |             |                                                              |
         |             |                                                      --------|---------
         |             |                                                      |                |
         |             |                                                      |                V
         |             |------------------------------------------------------------------->  mul
         |                                                                    |                |
         |                                                                    |                |
         |                                                                    V                |
         |-----------------------------------------------------------------> mul               |
                                                                              |                |
                                                                              |                V
                                                                              |              sub (beta) --> add --> [...]
                                                                              |                              ^
                                                                              |                              |
                                                                              |-------------------------------
  
  
      This pattern corresponds to either layer_norm or instance_norm.
  
      It is instance_norm if all of the following are true:
          - input is rank 4
          - axes of reduce_mean is [-2, -1]
          - gamma and beta are rank 1, after squeeze
  
      It is layer_norm if all of the following are true:
          - axes is either [-1] or [-1,-2] or [-1,-2,-3] and so on
          - rank of gamma and beta is equal to the length of the axes
      """
      for f_name, f in prog.functions.items():
          block_changed = True
          while block_changed:
              block_changed = fuse_layernorm_or_instancenorm_block(f)
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import absolute_import
  from coremltools.converters.mil.mil.passes.pass_registry import register_pass
  from coremltools.converters.mil.mil import Builder
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: fuse_layernorm_or_instancenorm(prog) without return types