def _can_fuse(expand_op: Operator, bmm_op: Operator) -> bool:
    """
    determine if expand_op and bmm_op can be fused
    """
    from aitemplate.compiler.ops.tensor.expand import (  # inner import to break circular import
        ExpandDimensionType,
    )

    expand_output = expand_op._attrs["outputs"][0]
    if expand_output._attrs["is_output"]:
        return False
    expand_inputs = expand_op._attrs["inputs"]
    expand_input_shape = expand_inputs[0]._attrs["shape"]
    expand_output_shape = expand_output._attrs["shape"]
    # not valid for bmm
    if len(expand_output_shape) != 3:
        return False
    if len(expand_input_shape) == 2:
        # In this case, we are expanding the batch dim
        assert (
            expand_input_shape[0] == expand_output_shape[1]
            and expand_input_shape[1] == expand_output_shape[2]
        ), f"invalid {expand_input_shape=} and {expand_output_shape=}"
        return True
    # not valid for bmm
    if len(expand_input_shape) != 3:
        return False
    if expand_op._attrs["dim_types"][0] != ExpandDimensionType.EXPAND_DIM:
        return False
    bmm_inputs = bmm_op._attrs["inputs"]
    bmm_a = bmm_inputs[0]
    bmm_b = bmm_inputs[1]
    if expand_output is bmm_a:
        return expand_output_shape[0] == bmm_a._attrs["shape"][0]
    if expand_output is bmm_b:
        return expand_output_shape[0] == bmm_b._attrs["shape"][0]
    return False
determine if expand_op and bmm_op can be fused
python/aitemplate/compiler/transform/fuse_expand_bmm.py