def compose_bmm(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    combine decomposed bmm (matmul)
    """
    modified = False
    # pattern replacement for aten_compose_mm_2d
    _LOGGER.info("compose_bmm: pattern matching for aten_compose_mm_2d...")
    aten_compose_mm_2d_replacement = replacement_pattern_abstract(aten_compose_mm_2d)
    res = replace_pattern(module, aten_compose_mm_2d, aten_compose_mm_2d_replacement)
    if len(res) > 0:
        modified = True
    # pattern replacement for aten_compose_bmm_2d
    _LOGGER.info("compose_bmm: pattern matching for aten_compose_bmm_3d...")

    def match_filter_aten_compose_bmm_2d(match, original_graph, pattern_graph):
        if len(match.placeholder_nodes[1].meta["val"].shape) == 2:
            return True
        else:
            return False

    aten_compose_bmm_2d_replacement = replacement_pattern_abstract(aten_compose_bmm_2d)
    res = replace_pattern(
        module,
        aten_compose_bmm_2d,
        aten_compose_bmm_2d_replacement,
        [match_filter_aten_compose_bmm_2d],
    )
    if len(res) > 0:
        modified = True
    # pattern replacement for aten_compose_bmm_3d
    _LOGGER.info("compose_bmm: pattern matching for aten_compose_bmm_2d...")

    def match_filter_aten_compose_bmm_3d(match, original_graph, pattern_graph):
        if len(match.placeholder_nodes[1].meta["val"].shape) == 3:
            return True
        else:
            return False

    aten_compose_bmm_3d_replacement = replacement_pattern_abstract(aten_compose_bmm_3d)
    res = replace_pattern(
        module,
        aten_compose_bmm_3d,
        aten_compose_bmm_3d_replacement,
        [match_filter_aten_compose_bmm_3d],
    )
    if len(res) > 0:
        modified = True

    return PassResult(module, modified)
combine decomposed bmm (matmul)
fx2ait/fx2ait/passes/lower_basic_pass_aten.py