def METHOD_NAME(
    example: Dict[Text, tf.Tensor], return_description: bool = False
) -> Optional[str]:
    """Print the keys in a example.
    Args:
        example: Dictionary keyed by strings with tensors as values.
        return_description: If `True`, returns the string description instead of
            printing it.
    Returns:
        String description if `return_description` is `True`, otherwise `None`.
    """
    desc = []
    key_length = max(len(k) for k in example.keys())
    for key, val in example.items():
        dtype = str(val.dtype) if isinstance(val.dtype, np.dtype) else repr(val.dtype)
        desc.append(
            f"{key.rjust(key_length)}: type={type(val).__name__}, "
            f"shape={val.shape}, "
            f"dtype={dtype}, "
            f"device={val.device if hasattr(val, 'device') else 'N/A'}"
        )
    desc = "\n".join(desc)
    if return_description:
        return desc
    else:
        print(desc)