[Shogi] move INIT_LEGAL_ACTION_MASK (#1275)

sotetsuk · Nov 1, 2024 · 7ed0cd8 · 7ed0cd8
1 parent 26a09bd
commit 7ed0cd8
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/pgx/_src/games/shogi.py b/pgx/_src/games/shogi.py
@@ -324,6 +324,15 @@ def f(i):
 
 ALL_SQ = jnp.arange(81)
 
+INIT_LEGAL_ACTION_MASK = np.zeros(81 * 27, dtype=jnp.bool_)
+# fmt: off
+ixs = np.int32([5, 7, 14, 23, 25, 32, 34, 41, 43, 50, 52, 59, 61, 68, 77, 79, 115, 124, 133, 142, 187, 196, 205, 214, 268, 277, 286, 295, 304, 331])
+# fmt: on
+INIT_LEGAL_ACTION_MASK[ixs] = True
+INIT_LEGAL_ACTION_MASK = jnp.bool_(INIT_LEGAL_ACTION_MASK)
+assert INIT_LEGAL_ACTION_MASK.shape == (81 * 27,)
+assert INIT_LEGAL_ACTION_MASK.sum() == 30
+
 
 class GameState(NamedTuple):
     turn: Array = jnp.int32(0)  # 0 or 1

diff --git a/pgx/shogi.py b/pgx/shogi.py
@@ -19,22 +19,12 @@
 import pgx.core as core
 from pgx._src.struct import dataclass
 from pgx._src.types import Array, PRNGKey
-from pgx._src.games.shogi import MAX_TERMINATION_STEPS, GameState, Game, _observe
+from pgx._src.games.shogi import MAX_TERMINATION_STEPS, GameState, Game, _observe, INIT_LEGAL_ACTION_MASK
 
 
 TRUE = jnp.bool_(True)
 FALSE = jnp.bool_(False)
 
-INIT_LEGAL_ACTION_MASK = jnp.zeros(81 * 27, dtype=jnp.bool_)
-# fmt: off
-ixs = [5, 7, 14, 23, 25, 32, 34, 41, 43, 50, 52, 59, 61, 68, 77, 79, 115, 124, 133, 142, 187, 196, 205, 214, 268, 277, 286, 295, 304, 331]
-# fmt: on
-for ix in ixs:
-    INIT_LEGAL_ACTION_MASK = INIT_LEGAL_ACTION_MASK.at[ix].set(True)
-assert INIT_LEGAL_ACTION_MASK.shape == (81 * 27,)
-assert INIT_LEGAL_ACTION_MASK.sum() == 30
-
-
 
 @dataclass
 class State(core.State):