def test_get_optimal_sequence(self):
"""Test ``get_optimal_sequence()``."""
_, generator, _ = set_random_seed(seed=42)
for shapes in _generate_shapes(generator=generator):
# get optimal sequence
opt_cost, opt_seq = get_optimal_sequence(*shapes)
# check correct cost
exp_opt_cost = estimate_cost_of_sequence(*(shapes[i] for i in opt_seq))
assert exp_opt_cost == opt_cost
# check optimality
for perm in itertools.permutations(list(range(len(shapes)))):
cost = estimate_cost_of_sequence(*(shapes[i] for i in perm))
assert cost >= opt_cost
def test_estimate_cost_of_add_sequence(self):
"""Test ``estimate_cost_of_add_sequence()``."""
_, generator, _ = set_random_seed(seed=42)
# create random array, estimate the costs of addition, and measure some execution times.
# then, compute correlation between the estimated cost, and the measured time.
data = []
for shapes in _generate_shapes(generator=generator):
arrays = [torch.empty(*shape) for shape in shapes]
cost = estimate_cost_of_sequence(*(a.shape for a in arrays))
n_samples, time = timeit.Timer(stmt="sum(arrays)", globals=dict(arrays=arrays)).autorange()
consumption = time / n_samples
data.append((cost, consumption))
a = numpy.asarray(data)
# check for strong correlation between estimated costs and measured execution time
assert (numpy.corrcoef(x=a[:, 0], y=a[:, 1])[0, 1]) > 0.8