code: def try_to_transform(onehot_op, block):
      root_var = onehot_op.indices
  
      # check that the output of the onehot op is not a block output
      if onehot_op.outputs[0] in block.outputs:
          return False
  
      # check that onehot op has axis=-1, on_value=1 and off_value=0
      # and constant one_hot_vector_size
      axis = onehot_op.axis.val
      if axis is None:
          return False
      if onehot_op.indices.shape is None:
          return False
      rank = len(onehot_op.indices.shape)
      if axis >= 0:
          axis -= rank
      if axis != -1:
          return False
      if not _check_var_scalar_value(onehot_op.on_value, 1):
          return False
      if not _check_var_scalar_value(onehot_op.off_value, 0):
          return False
      if onehot_op.one_hot_vector_size.val is None:
          return False
  
      # checks for the following matmul op
      if not _check_child_op_type(onehot_op, "matmul"):
          return False
      matmul_op = list(onehot_op.outputs[0].child_ops)[0]
      if matmul_op.x != onehot_op.outputs[0]:
          return False
      if matmul_op.transpose_x.val or matmul_op.transpose_y.val:
          return False
      W_var = matmul_op.y
      if W_var.val is None:
          return False
      if len(W_var.val.shape) != 2:
          return False
  
      # remove onehot and matmul and replace with gather op
      out_name = matmul_op.outputs[0].name
      x = mb.gather(
          x=W_var,
          indices=root_var,
          axis=0,
          name=out_name,
          before_op=matmul_op,
      )
  
      matmul_op.enclosing_block.replace_uses_of_var_after_op(
          anchor_op=matmul_op, old_var=matmul_op.outputs[0], new_var=x
      )
      # Remove all the ops at once
      block.remove_ops([onehot_op, matmul_op])
      return True
  
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 3 for calling a global function with example usage: try_to_transform(onehot_op, block) and returns: True, False, False, False, False, False, False, False, False, False, False, False, False