test_nan.py

import struct

def check_bf16_nan(filename):
    """
    Checks if the given BF16 binary file contains any NaN values.
    Assumes the file is in big-endian byte order.
    Returns True if any NaN is found, False otherwise.
    """
    has_nan = False
    with open(filename, 'rb') as f:
        while True:
            chunk = f.read(2)
            if not chunk:
                break  # End of file
            if len(chunk) != 2:
                print("Warning: Incomplete last chunk (expected 2 bytes).")
                break
            # Unpack as big-endian 16-bit unsigned integer
            value = struct.unpack('>H', chunk)[0]
            # Extract exponent (bits 14-7) and mantissa (bits 6-0)
            exponent = (value >> 7) & 0xFF
            mantissa = value & 0x7F
            if exponent == 0xFF and mantissa != 0:
                has_nan = True
                print(f"NaN detected at byte offset {f.tell() - 2}: value=0x{value:04X}")
    return has_nan

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("Usage: python check_bf16_nan.py <filename>")
        sys.exit(1)
    filename = sys.argv[1]
    if check_bf16_nan(filename):
        print("The file contains NaNs.")
        sys.exit(2)
    else:
        print("The file does not contain NaNs.")
        sys.exit(0)