[Shogi] Add and fix buggy test samples

pgx/shogi.py

@@ -929,9 +929,7 @@ def _sfen_to_state(sfen):
     else:
         s_turn = jnp.int8(1)
     s_hand = jnp.zeros(14, dtype=jnp.int8)
-    if hand == "-":
-        s_hand = jnp.reshape(s_hand, (2, 7))
-    else:
+    if hand != "-":
         num_piece = 1
         for char in hand:
             if char.isdigit():
@@ -942,7 +940,7 @@ def _sfen_to_state(sfen):
     return State(
         turn=s_turn,
         piece_board=jnp.rot90(piece_board.reshape((9, 9)), k=1).flatten(),
-        hand=s_hand,
+        hand=jnp.reshape(s_hand, (2, 7)),
     )
 
 

tests/test_shogi.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 
 from pgx.shogi import *
-from pgx.shogi import _init, _step, _step_move, _step_drop, _flip, _apply_effects, _legal_actions, _rotate, _to_direction
+from pgx.shogi import _init, _step, _step_move, _step_drop, _flip, _apply_effects, _legal_actions, _rotate, _to_direction, _sfen_to_state
 
 
 # check visualization results by image preview plugins
@@ -373,3 +373,29 @@ def test_step():
     s = step(s, 3 * 81 + xy2i(3, 8))
     visualize(s, "tests/assets/shogi/step_003.svg")
     assert not s.legal_action_mask[3 * 81 + xy2i(3, 8)]
+
+
+# 今やると落ちる
+def test_legal_action_mask():
+    # 歩以外の持ち駒に対しての二歩判定回避
+    sfen = "9/9/9/9/9/9/PPPPPPPPP/9/9 b NLP 1"
+    legal_action = _to_direction(_legal_actions(_sfen_to_state(sfen)))
+    # 歩は二歩になるので打てない
+    assert (legal_action[81*20:81*21] == False).all()
+    # 香車は1列目と6列目（歩がいる）二は打てない
+    # 二列目には打てる
+    assert (legal_action[81*21+1:81*22:9] == True).all()
+    assert (legal_action[81*21:81*22:9] == False).all()
+    assert (legal_action[81*21+5:81*22:9] == False).all()
+    # 桂馬は1.2列目に打てない
+    # 三列目には打てる
+    assert (legal_action[81*22+1:81*23:9] == False).all()
+    assert (legal_action[81*22:81*23:9] == False).all()
+    assert (legal_action[81*21+5:81*22:9] == False).all()
+
+    # 成駒のpromotion判定
+    sfen = "9/2+B1G1+P2/9/9/9/9/9/9/9 b - 1"
+    legal_action = _to_direction(_legal_actions(_sfen_to_state(sfen)))
+    legal_action = jnp.where(legal_action, 1, 0)
+    # promotionは生成されてたらダメ
+    assert jnp.all(legal_action[810:] == 0)