code: def fuse_onehot_matmul_to_gather(prog):
      """
      Detect if onehot (axis=-1, on_value=1, off_value=0) is followed by a matmul op (no bias),
      then they can be replaced by a gather op.
  
      Input:
          %2 = one_hot(%1, on_value=1, off_value=0, axis=-1)
          %3 = const() # rank 2
          %4  = matmul(%2, %3)
  
      Output:
          %4 = gather(%3, %2, axis=0)
  
      """
      for f_name, f in prog.functions.items():
          block_changed = True
          while block_changed:
              block_changed = fuse_onehot_matmul_to_gather_block(f)
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.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 5 for calling a global function with example usage: fuse_onehot_matmul_to_gather(prog) without return types