def softmax_gen_function(func_attrs: Dict[str, Any]) -> str:
    """Generate function body.

    Parameters
    ----------
    func_attrs : Dict
        Operation attributes.

    Returns
    -------
    str
        The rendered template of generated function body.
    """
    dim = func_attrs["dim"]
    shapes = func_attrs["inputs"][0]._attrs["shape"]
    rank = len(shapes)

    assert (
        dim == rank - 1
    ), f"rocm softmax only supports dim == rank - 1, dim={dim}, rank={rank}"

    assert isinstance(
        shapes[dim], IntImm
    ), "softmax requires reduction dim to be static"
    return norm_common.gen_function(
        func_attrs, EXEC_TEMPLATE, EXTRA_HEADERS, get_func_signature
    )
Generate function body.

    Parameters
    ----------
    func_attrs : Dict
        Operation attributes.

    Returns
    -------
    str
        The rendered template of generated function body.
    
python/aitemplate/backend/rocm/normalization/softmax.py