code: def fuse_elementwise_to_batchnorm(prog):
      """
      Fold mul + add into a batch norm,
      if the const feeding into the mul/add is of shape (1,C,1,1) or (C,1,1)
      and input to mul is of rank 4.
  
      Given:
               [Const]   [Const]
                  |         |
                  V         V
      [...] --> [Mul] --> [Add] --> [...]
  
      That is,
  
          %2 = op1(%1)
          %3 = mul(%2, constant)
          %4 = add(%3, constant)
          %5 = op2(%4)
          ...
  
      Result:
  
      [...] --> [BatchNorm] --> [...]
  
      That is,
          %2 = op1(%1)
          %4 = batchnorm(%2)
          %5 = op2(%4)
          ...
  
      """
      for f_name, f in prog.functions.items():
          block_changed = True
          while block_changed:
              block_changed = fuse_elementwise_to_batchnorm_block(f)
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/mil/passes/elementwise_batchnorm_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 6 for calling a global function with example usage: fuse_elementwise_to_batchnorm(prog) without return types