code: def test_onehot_matmul_to_gather_fusion(rank):
      """
      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)
      """
      rank4_shape = (10, 3, 6, 7)
      input_shape = rank4_shape[-rank:]
      vocab_size = 15
      embedding_size = 12
  
      @mb.program(
          input_specs=[
              mb.TensorSpec(shape=input_shape, dtype=types.int32)
          ]
      )
      def prog(x):
          x = mb.one_hot(
              indices=x,
              on_value=1,
              off_value=0,
              axis=-1,
              one_hot_vector_size=vocab_size,
          )
          x = mb.matmul(
              x=x, y=np.random.rand(vocab_size, embedding_size)
          )
          return x
  
      prev_prog, prev_block, block = apply_pass_and_basic_check(
          prog, "common::fuse_onehot_matmul_to_gather"
      )
      assert get_op_types_in_program(prev_prog) == ["one_hot", "matmul"]
      assert get_op_types_in_program(prog) == ["gather"]
      assert_model_is_valid(
          prog,
          {"x": input_shape},
          expected_output_shapes={
              block.outputs[0].name: input_shape + (embedding_size,)
          },
      )
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/mil/passes/test_passes.py
imports: from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.testing_utils import assert_op_count_match, assert_model_is_valid, assert_same_output_names, get_op_types_in_program, apply_pass_and_basic_check
  from coremltools.converters.mil.mil import Symbol, types
  from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
  import copy
  import pytest
  import numpy
label_desc: Use image node_id 11 for calling a global function with example usage: test_onehot_matmul_to_gather_fusion(rank) and returns: x