Function generate_random_array

python/tvm/testing/utils.py:1205–1247

Generate a random array by generating random bits and casting to the target dtype. Supported dtypes: - "int8", "uint8", "float16", "float32", "bfloat16", "float8_e4m3fn", "float8_e5m2"

(dtype: str, shape: tuple)

def generate_random_array(dtype: str, shape: tuple) -> np.ndarray:
    """
    Generate a random array by generating random bits and casting to the target dtype.

    Supported dtypes:
      - "int8", "uint8", "float16", "float32", "bfloat16", "float8_e4m3fn", "float8_e5m2"
    """
    try:
        np_dtype = np_dtype_from_str(dtype)

    except TypeError:
        raise ValueError("Provided dtype is not a valid numpy dtype.")

    # Determine the bit length for this dtype.
    bit_length = np_dtype.itemsize * 8

    # Choose an appropriate unsigned container type.
    if bit_length <= 8:
        container = np.uint8
    elif bit_length <= 16:
        container = np.uint16
    elif bit_length <= 32:
        container = np.uint32
    elif bit_length <= 64:
        container = np.uint64
    else:
        raise ValueError(f"Unsupported dtype bit length: {bit_length}")

    # Generate random integers in the full range of the bit length.
    random_ints = np.random.randint(0, 2**bit_length, size=shape, dtype=container)
    # Reinterpret the bit pattern as the desired dtype.
    res = random_ints.view(np_dtype)
    with np.errstate(invalid="ignore"):
        invalid_indices = np.where(~np.isfinite(res))
    for idx in zip(*invalid_indices):
        while True:
            with np.errstate(invalid="ignore"):
                if np.isfinite(res[idx]):
                    break
            # Generate a new random value for this specific position
            new_random_int = np.random.randint(0, 2**bit_length, size=1, dtype=container)
            res[idx] = new_random_int.view(np_dtype)[0]
    return res