def get_gen_function(a_layout, b_layout, c_layout):
    """
    Return gen_function for given layouts of A, B, and bias.
    """

    def gen_function(
        func_attrs,
        exec_cond_template,
        dim_info_dict,
    ):
        problem_args = _get_problem_args(a_layout, b_layout, c_layout)

        default_mm_info = bmm_common.get_default_problem_info(
            problem_args,
            alpha_value=func_attrs.get("alpha", 1),
        )
        (
            problem_args,
            problem_args_cutlass_3x,
            input_addr_calculator,
            output_addr_calculator,
        ) = bmm_common.make_function_strided_args(
            func_attrs=func_attrs,
            dim_info_dict=dim_info_dict,
            default_mm_info=default_mm_info,
            is_permute=False,
        )

        return bmm_common.gen_function(
            func_attrs=func_attrs,
            exec_cond_template=exec_cond_template,
            problem_args=problem_args,
            problem_args_cutlass_3x=problem_args_cutlass_3x,
            dim_info_dict=dim_info_dict,
            input_addr_calculator=input_addr_calculator,
            output_addr_calculator=output_addr_calculator,
        )

    return gen_function
Return gen_function for given layouts of A, B, and bias.
python/aitemplate/backend/cuda/gemm_universal/bmm_xxx.py