change base dtype from int8 to int16 to avoid overflow when cube size is >= 7

src/rubik/action.py

@@ -15,7 +15,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 0, 1],
                 [0, 0, -1, 0],
             ],
-            dtype=torch.int8,
+            dtype=torch.int16,
         ),
         # rot about Y: X -> Z
         torch.tensor(
@@ -25,7 +25,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 1, 0],
                 [0, 1, 0, 0],
             ],
-            dtype=torch.int8,
+            dtype=torch.int16,
         ),
         # rot about Z: Y -> X
         torch.tensor(
@@ -35,7 +35,7 @@ POS_ROTATIONS = torch.stack(
                 [0, -1, 0, 0],
                 [0, 0, 0, 1],
             ],
-            dtype=torch.int8,
+            dtype=torch.int16,
         ),
     ]
 )
@@ -46,7 +46,7 @@ POS_SHIFTS = torch.tensor(
         [0, 1, 0, 0],
         [0, 0, 1, 0],
     ],
-    dtype=torch.int8,
+    dtype=torch.int16,
 )
 
 
@@ -74,7 +74,7 @@ def build_actions_tensor(size: int) -> torch.Tensor:
             for inverse in range(2)
         ],
         dim=0,
-    ).sum(dim=0, dtype=torch.int8)
+    ).sum(dim=0, dtype=torch.int16)
 
 
 def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch.Tensor:
@@ -87,7 +87,7 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     length = 6 * (size**2)
 
     # extract faces impacted by the move
-    indices = tensor.indices().to(dtype=torch.int8)  # size = (4, length)
+    indices = tensor.indices().to(dtype=torch.int16)  # size = (4, length)
     changes = (indices[axis + 1] == slice).nonzero().reshape(-1)  # size = (n,), n < length
     extract = indices[:, changes]  # size = (4, n)
 
@@ -97,7 +97,7 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     rotated = rotated + offsets  # size = (4, n)
 
     # apply face rotation
-    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int8) @ FACE_ROTATIONS[axis]).argmax(dim=-1)
+    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int16) @ FACE_ROTATIONS[axis]).argmax(dim=-1)
 
     # from this point on, convert rotation into a position-based permutation of colors
     (inputs, outputs) = (rotated, extract) if bool(inverse) else (extract, rotated)
@@ -116,11 +116,11 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     # convert permutation dict into sparse tensor
     perm_indices = torch.tensor(
         [[axis] * length, [slice] * length, [inverse] * length, list(total_perm.keys()), list(total_perm.values())],
-        dtype=torch.int8,
+        dtype=torch.int16,
     )
-    perm_values = torch.tensor([1] * length, dtype=torch.int8)
+    perm_values = torch.tensor([1] * length, dtype=torch.int16)
     perm_size = (3, size, 2, length, length)
-    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int8)
+    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int16)
 
 
 def parse_action_str(move: str) -> tuple[int, ...]:

src/rubik/cube.py

@@ -28,8 +28,8 @@ class Cube:
             cube = Cube(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
         """
         tensor = build_cube_tensor(colors, size)
-        self.coordinates = tensor.indices().transpose(0, 1).to(torch.int8)
-        self.state = F.one_hot(tensor.values().long()).to(torch.int8)
+        self.coordinates = tensor.indices().transpose(0, 1).to(torch.int16)
+        self.state = F.one_hot(tensor.values().long()).to(torch.int16)
         self.actions = build_actions_tensor(size)
         self.history: list[list[int]] = []
         self.colors = colors
@@ -37,7 +37,7 @@ class Cube:
 
     def to(self, device: str | torch.device) -> "Cube":
         device = torch.device(device)
-        dtype = torch.int8 if device == torch.device("cpu") else torch.float32
+        dtype = torch.int16 if device == torch.device("cpu") else torch.float32
         self.coordinates = self.coordinates.to(device=device, dtype=dtype)
         self.state = self.state.to(device=device, dtype=dtype)
         self.actions = self.actions.to(device=device, dtype=dtype)
@@ -84,7 +84,7 @@ class Cube:
         """
         actions = parse_actions_str(moves)
         tensors = [self.actions[*action].to(torch.float32) for action in actions]
-        result = reduce(lambda A, B: A @ B, tensors).to(torch.int8)
+        result = reduce(lambda A, B: A @ B, tensors).to(torch.int16)
         return dict(result.indices().transpose(0, 1).tolist())
 
     def solve(self, policy: str) -> None:
@@ -97,5 +97,5 @@ class Cube:
         """
         Compute a string representation of a cube.
         """
-        state = self.state.argmax(dim=-1).to(device="cpu", dtype=torch.int8)
+        state = self.state.argmax(dim=-1).to(device="cpu", dtype=torch.int16)
         return stringify(state, self.colors, self.size)

src/rubik/tensor_utils.py

@@ -10,7 +10,7 @@ def build_cube_tensor(colors: list[str], size: int) -> torch.Tensor:
 
     # build dense tensor filled with colors
     n = size - 1
-    tensor = torch.zeros([6, size, size, size], dtype=torch.int8)
+    tensor = torch.zeros([6, size, size, size], dtype=torch.int16)
     tensor[0, :, :, n] = 1  # up
     tensor[1, 0, :, :] = 2  # left
     tensor[2, :, n, :] = 3  # front
@@ -26,6 +26,6 @@ def build_permutation_matrix(size: int, perm: str) -> torch.Tensor:
     """
     perm_list = [int(p) for p in (perm + perm[0])]
     perm_dict = {perm_list[i]: perm_list[i + 1] for i in range(len(perm))}
-    indices = torch.tensor([list(range(size)), [(perm_dict.get(i, i)) for i in range(size)]], dtype=torch.int8)
-    values = torch.tensor([1] * size, dtype=torch.int8)
-    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int8)
+    indices = torch.tensor([list(range(size)), [(perm_dict.get(i, i)) for i in range(size)]], dtype=torch.int16)
+    values = torch.tensor([1] * size, dtype=torch.int16)
+    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int16)