|
import torch |
|
|
|
from sudoku.models import SudokuNet |
|
|
|
|
|
def test_same_output_under_rotation(): |
|
model = SudokuNet() |
|
arr1 = torch.zeros((1, 2, 9, 9, 9)) |
|
arr1[0, 0, 1, 2, 3] = 1 |
|
|
|
output_1 = model.forward(arr1.view(1, 2, 9 * 9 * 9)) |
|
assert output_1.shape == (1, 2, 9 * 9 * 9), output_1 |
|
arr2 = torch.zeros((1, 2, 9, 9, 9)) |
|
arr2[0, 0, 2, 3, 4] = 1 |
|
|
|
output_2 = model.forward(arr2.view(1, 2, 9 * 9 * 9)) |
|
assert ( |
|
output_1.view(1, 2, 9, 9, 9)[0, 0, 1, 2, 3] |
|
== output_2.view(1, 2, 9, 9, 9)[0, 0, 2, 3, 4] |
|
) |
|
assert ( |
|
output_1.view(1, 2, 9, 9, 9)[0, 0, 1, 2, 4] |
|
== output_2.view(1, 2, 9, 9, 9)[0, 0, 2, 3, 6] |
|
) |
|
assert ( |
|
output_1.view(1, 2, 9, 9, 9)[0, 1, 1, 2, 4] |
|
== output_2.view(1, 2, 9, 9, 9)[0, 1, 2, 3, 6] |
|
) |
|
assert ( |
|
output_1.view(1, 2, 9, 9, 9)[0, 1, 2, 2, 4] |
|
== output_2.view(1, 2, 9, 9, 9)[0, 1, 1, 3, 6] |
|
) |
|
assert ( |
|
output_1.view(1, 2, 9, 9, 9)[0, 1, 2, 3, 4] |
|
== output_2.view(1, 2, 9, 9, 9)[0, 1, 1, 2, 6] |
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
