code: def try_to_transform(mul_op, add_op, block):
      non_const_input_mul = (
          mul_op.x if mul_op.x.val is None else mul_op.y
      )
      if non_const_input_mul.rank != 4:
          return False
  
      gamma = _find_const_input_val(mul_op)
      beta = _find_const_input_val(add_op)
      if gamma is None or beta is None:
          return False
  
      if not (
          isinstance(gamma, np.ndarray) and isinstance(beta, np.ndarray)
      ):
          return False
  
      # check that gamma and beta have shape (1,C,1,1) or (C,1,1)
      # that is they are doing vector addition on the axis=-3, which is what the
      # batchnorm layer does (batchnorm layer only works on rank 4 input tensors)
      if not (_check_shape(gamma) and _check_shape(beta)):
          return False
  
      C = gamma.shape[-3]
      if C == 1:
          return False
  
      out_name = add_op.outputs[0].name
      x = mb.batch_norm(
          x=non_const_input_mul,
          mean=np.zeros((C,), np.float32),
          variance=np.ones((C,), np.float32),
          gamma=np.squeeze(gamma),
          beta=np.squeeze(beta),
          name=out_name,
          before_op=mul_op,
      )
  
      add_op.enclosing_block.replace_uses_of_var_after_op(
          anchor_op=add_op, old_var=add_op.outputs[0], new_var=x
      )
      # Remove all the ops at once
      block.remove_ops([mul_op, add_op])
      return True
  
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 4 for calling a global function with example usage: try_to_transform(mul_op, add_op, block) and returns: True, False, False, False, False, False