def _filter_gemm_op(op: Operator) -> bool:
    """Function to filter out bad gemm candidates for group_gemm."""
    if op._attrs["alpha"] != 1.0:
        return False

    # 2D input
    input = op._attrs["inputs"][0]
    if input._rank() != 2:
        return False

    # dynamic dim check
    if not all_static_dimensions(input.shape()):
        return False

    return _get_ab_alignment(op) > 1
Function to filter out bad gemm candidates for group_gemm.
python/aitemplate/compiler/transform/fuse_group_ops.py