Upload attn_map.py with huggingface_hub

attn_map.py

@@ -0,0 +1,116 @@
+import numpy as np
+import torch
+
+move = np.arange(1, 8)
+
+diag = np.array([
+    move    + move*8,
+    move    - move*8,
+    move*-1 - move*8,
+    move*-1 + move*8
+])
+
+orthog = np.array([
+    move,
+    move*-8,
+    move*-1,
+    move*8
+])
+
+knight = np.array([
+    [2 + 1*8],
+    [2 - 1*8],
+    [1 - 2*8],
+    [-1 - 2*8],
+    [-2 - 1*8],
+    [-2 + 1*8],
+    [-1 + 2*8],
+    [1 + 2*8]
+])
+
+promos = np.array([2*8, 3*8, 4*8])
+pawn_promotion = np.array([
+    -1 + promos,
+    0 + promos,
+    1 + promos
+])
+
+def make_map():
+    """theoretically possible put-down squares (numpy array) for each pick-up square (list element).
+    squares are [0, 1, ..., 63] for [a1, b1, ..., h8]. squares after 63 are for promotion squares.
+    each successive "row" beyond 63 (ie. 64:72, 72:80, 80:88) are for over-promotions to queen, rook, and bishop;
+    respectively. a pawn traverse to row 56:64 signifies a "default" promotion to a knight."""
+    traversable = []
+    for i in range(8):
+        for j in range(8):
+            sq = (8*i + j)
+            traversable.append(
+                sq +
+                np.sort(
+                    np.int32(
+                        np.concatenate((
+                            orthog[0][:7-j], orthog[2][:j], orthog[1][:i], orthog[3][:7-i],
+                            diag[0][:np.min((7-i, 7-j))], diag[3][:np.min((7-i, j))],
+                            diag[1][:np.min((i, 7-j))], diag[2][:np.min((i, j))],
+                            knight[0] if i < 7 and j < 6 else [], knight[1] if i > 0 and j < 6 else [],
+                            knight[2] if i > 1 and j < 7 else [], knight[3] if i > 1 and j > 0 else [],
+                            knight[4] if i > 0 and j > 1 else [], knight[5] if i < 7 and j > 1 else [],
+                            knight[6] if i < 6 and j > 0 else [], knight[7] if i < 6 and j < 7 else [],
+                            pawn_promotion[0] if i == 6 and j > 0 else [],
+                            pawn_promotion[1] if i == 6           else [],
+                            pawn_promotion[2] if i == 6 and j < 7 else [],
+                        ))
+                    )
+                )
+            )
+    z = np.zeros((64*64+8*24, 1858), dtype=np.int32)
+    apm_out = np.zeros((1858,), dtype=np.int32)
+    apm_in = np.zeros((64*64+8*24), dtype=np.int32)
+    # first loop for standard moves (for i in 0:1858, stride by 1)
+    i = 0
+    for pickup_index, putdown_indices in enumerate(traversable):
+        for putdown_index in putdown_indices:
+            if putdown_index < 64:
+                du_idx = putdown_index + (64*pickup_index)
+                z[du_idx, i] = 1
+                apm_out[i] = du_idx
+                apm_in[du_idx] = i
+                i += 1
+    # second loop for promotions (for i in 1792:1858, stride by ls[j])
+    j = 0
+    j1 = np.array([3, -2, 3, -2, 3])
+    j2 = np.array([3, 3, -5, 3, 3, -5, 3, 3, 1])
+    ls = np.append(j1, 1)
+    for k in range(6):
+        ls = np.append(ls, j2)
+    ls = np.append(ls, j1)
+    ls = np.append(ls, 0)
+    for pickup_index, putdown_indices in enumerate(traversable):
+        for putdown_index in putdown_indices:
+            if putdown_index >= 64:
+                pickup_file = pickup_index % 8
+                promotion_file = putdown_index % 8
+                promotion_rank = (putdown_index // 8) - 8
+                du_idx = 4096 + pickup_file*24 + (promotion_file*3+promotion_rank)
+                z[du_idx, i] = 1
+                apm_out[i] = du_idx
+                apm_in[du_idx] = i
+                i += ls[j]
+                j += 1
+
+    return z, apm_out, apm_in
+
+apm_map, apm_out, apm_in = make_map()
+
+def set_zero_sum(x):
+    x = x + (1 - torch.sum(x, dim=1, keepdim=True)) * (1.0 / 64)
+    return x
+
+def get_up_down(moves):
+    apm_map_tensor = torch.from_numpy(apm_map)
+    out = torch.matmul(moves, apm_map_tensor.T.float())
+    out = out[..., :64*64]
+    out = out.view(-1, 64, 64)
+    pu = set_zero_sum(torch.sum(out, dim=-1))
+    pd = set_zero_sum(torch.sum(out, dim=-2))
+    return pu, pd