def dtype_to_enumerator(dtype: str) -> str:
    """Returns the string representation of the AITemplateDtype enum
    (defined in model_interface.h) for the given dtype str.

    Parameters
    ----------
    dtype: str
        A data type string.

    Returns
    ----------
    str
        the AITemplateDtype enum string representation.
    """

    def _impl(dtype):
        if dtype == "float16":
            return "kHalf"
        elif dtype == "float32" or dtype == "float":
            return "kFloat"
        elif dtype == "int32" or dtype == "int":
            return "kInt"
        elif dtype == "int64":
            return "kLong"
        elif dtype == "bool":
            return "kBool"
        elif dtype == "bfloat16":
            return "kBFloat16"
        else:
            raise AssertionError(f"unknown dtype {dtype}")

    return f"AITemplateDtype::{_impl(dtype)}"
Returns the string representation of the AITemplateDtype enum
    (defined in model_interface.h) for the given dtype str.

    Parameters
    ----------
    dtype: str
        A data type string.

    Returns
    ----------
    str
        the AITemplateDtype enum string representation.
    
python/aitemplate/compiler/dtype.py