def test_fuse_onehot_matmul_to_gather(self, backend, rank, opset_version):
        """
        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)],
            opset_version=opset_version,
        )
        def prog(x):
            x = mb.one_hot(
                indices=x, on_value=1.0, off_value=0.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"]
        if opset_version == ct.target.iOS17:
            # Several ops added to make sure indices in iOS17 gather is non-negative.
            assert get_op_types_in_program(prog) == [
                "greater_equal",
                "shape",
                "slice_by_index",
                "add",
                "select",
                "gather",
            ]
        else:
            assert get_op_types_in_program(prog) == ["gather"]

        if opset_version == ct.target.iOS17:
            if backend[0] != "mlprogram" or _macos_version() < (14, 0):
                pytest.skip("IOS17 target available only on macOS 14+ with mlprogram.")

        assert_model_is_valid(
            prog,
            {"x": input_shape},
            backend=backend,
            expected_output_shapes={block.outputs[0].name: input_shape + (embedding_size,)},
            minimum_deployment_target=opset_version,
        )

        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)
        
coremltools/converters/mil/mil/passes/tests/test_passes.py