argmax_test.py

"""
Create tests for the argmax

Usage:
python3 argmax_test.py <n_inputs> <bit_width> <n_tests>

Output format:
The .dat file has the following format (example with 4 inputs, i_n n=1,2,3,4):
<i_1> <i_2> <i_3> <i_4> <argmax>
"""

import argparse
import random
from math import ceil, log2

import numpy as np
from test_module import write_dat  # noqa: E402

random.seed(7357)

parser = argparse.ArgumentParser(description="Testing script for argmax")
parser.add_argument("-inp", help="Number of inputs", type=int)
parser.add_argument("-bw", help="Bit width", type=int)
parser.add_argument("-tests", help="Number of tests", type=int)
args = parser.parse_args()

# Valid range of the inputs
RANGE = int(2 ** (args.bw - 1))

# Delay of the block
delay = ceil(log2(args.inp)) + 2

# Random test inputs
test_inputs = [random.choices(range(-RANGE, RANGE - 1), k=args.inp) for _ in range(args.tests)]

# Delay the inputs
for _ in range(delay):
    test_inputs.append([0 for _ in range(args.inp)])

# Delay the results
zero_pad = [0] * delay
maxs = [np.argmax(np.array(vec)) for vec in test_inputs]
results = zero_pad + maxs

# Write test cases to file
write_dat("argmax.dat", args.tests + delay, test_inputs, results)