from typing import Dict
import torch
from typing import List


def find_node_fp32(graph: Dict[str, str], output_nodes: Dict[str, torch.Tensor]) -> List[str]:
    """
    Identify out of range values in node outputs.
    :param graph: graph as adjency nodes dict
    :param output_nodes: output of each node
    :return: list of nodes producing outputs outside fp16 tensor
    """
    keep_fp32 = list()
    min_float16 = torch.finfo(torch.float16).min
    max_float16 = torch.finfo(torch.float16).max
    resolution = 5.96e-08  # torch.finfo(torch.float16).eps  # minimum value that can be represented by FP16
    for k, tensor in output_nodes.items():
        if tensor.dtype != torch.float32:
            continue
        # out of FP16 range
        if (
            torch.any(tensor > max_float16)
            or torch.any(tensor < min_float16)
            or (torch.any((tensor < resolution) & (tensor > -resolution) & (tensor != 0)))  # limited memory footprint
        ):
            keep_fp32.append(graph[k])
    return keep_fp32