Repository: sotetsuk/pgx

Issue:
[Shogi] Enable jitting
85個関数がある


Discussion:

Resolving patch:
diff --git a/pgx/shogi.py b/pgx/shogi.py
--- a/pgx/shogi.py
+++ b/pgx/shogi.py
@@ -1,9 +1,20 @@
-import copy
+from functools import partial
 from typing import Tuple
 
+import jax
 import jax.numpy as jnp
+from flax.serialization import from_bytes
 from flax.struct import dataclass
 
+TRUE = jnp.bool_(True)
+FALSE = jnp.bool_(False)
+
+
+ACTION_SIZE = 2754
+INIT_BOARD_BYTES = b"\xc8$\xc1\x01\x93\x92\x1dQ\xa5int32\xc5$\xb4\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+INIT_BOARD = from_bytes(jnp.zeros((29, 81), dtype=jnp.int32), INIT_BOARD_BYTES)
+INIT_BOARD = jnp.array(INIT_BOARD)
+
 
 # 指し手のdataclass
 @dataclass
@@ -43,136 +54,68 @@ class ShogiState:
 
 
 def init() -> ShogiState:
-    board = _make_init_board()
+    board = INIT_BOARD
     state = ShogiState(board=board, hand=jnp.zeros(14, dtype=jnp.int32))  # type: ignore
     return _init_legal_actions(state)
 
 
 def step(state: ShogiState, action: int) -> Tuple[ShogiState, int, bool]:
     # state, 勝敗判定,終了判定を返す
-    s = copy.deepcopy(state)
-    legal_actions = _legal_actions(s)
-    _action = _dlaction_to_action(action, s)
+    legal_actions = _legal_actions(state)
+    _action = _dlaction_to_action(action, state)
+
     # actionのfromが盤外の場合は非合法手なので負け
-    if not _is_in_board(_action.from_):
-        print("an illegal action")
-        return s, _turn_to_reward(_another_color(s)), True
-    # legal_actionsにactionがない場合、そのactionは非合法手
-    if legal_actions[action] == 0:
-        print("an illegal action2")
-        return s, _turn_to_reward(_another_color(s)), True
+    is_oob = ~_is_in_board(_action.from_)
+    is_illegal = legal_actions[action] == 0
+    state1 = state
+
     # 合法手の場合
-    # 駒打ち
-    if _action.is_drop:
-        s = _update_legal_drop_actions(s, _action)
-        s = _drop(s, _action)
-        print("drop: piece =", _action.piece, ", to =", _action.to)
-    # 駒の移動
-    else:
-        s = _update_legal_move_actions(s, _action)
-        s = _move(s, _action)
-        print("move: piece =", _action.piece, ", to =", _action.to)
-    # 王手がかかったままの場合、王手放置またｈ自殺手で負け
-    cn, cnp, cf, cfp = _is_check(s)
-    if cn + cf != 0:
-        print("check is remained")
-        return s, _turn_to_reward(_another_color(s)), True
-    # その他の反則
-    if _is_double_pawn(s):
-        print("two pawns in the same file")
-        return s, _turn_to_reward(_another_color(s)), True
-    if _is_stuck(s):
-        print("some pieces are stuck")
-        return s, _turn_to_reward(_another_color(s)), True
-    s = s.replace(turn=_another_color(s))  # type: ignore
-    # 相手に合法手がない場合→詰み
-    if _is_mate(s):
-        # actionのis_dropがTrueかつpieceが歩の場合、打ち歩詰めで負け
-        if _action.is_drop and (_action.piece == 1 or _action.piece == 15):
-            print("mate by dropped pawn")
-            return s, _turn_to_reward(s.turn), True
-        # そうでなければ普通の詰みで勝ち
-        else:
-            print("mate")
-            return s, _turn_to_reward(_another_color(s)), True
-    else:
-        return s, 0, False
+    state = jax.lax.cond(
+        _action.is_drop,
+        lambda: _drop(_update_legal_drop_actions(state, _action), _action),
+        lambda: _move(_update_legal_move_actions(state, _action), _action),
+    )
+
+    # 王手がかかったままの場合、王手放置また自殺手で負け
+    cn, cnp, cf, cfp = _is_check(state)
+    is_check = cn + cf != 0
+    is_double_pawn = _is_double_pawn(state)
+    is_stuck = _is_stuck(state)
+    state2 = state
 
+    state = state.replace(turn=_another_color(state))  # type: ignore
 
-# turnから報酬計算
-def _turn_to_reward(turn: int) -> int:
-    if turn == 0:
-        return 1
-    else:
-        return -1
-
-
-# 初期盤面生成
-def _make_init_board() -> jnp.ndarray:
-    array = jnp.zeros((29, 81), dtype=jnp.int32)
-    for i in range(81):
-        if i % 9 == 6:
-            # 先手歩
-            p = 1
-        elif i == 8 or i == 80:
-            # 先手香車
-            p = 2
-        elif i == 17 or i == 71:
-            # 先手桂馬
-            p = 3
-        elif i == 26 or i == 62:
-            # 先手銀
-            p = 4
-        elif i == 70:
-            # 先手角
-            p = 5
-        elif i == 16:
-            # 先手飛車
-            p = 6
-        elif i == 35 or i == 53:
-            # 先手金
-            p = 7
-        elif i == 44:
-            # 先手玉
-            p = 8
-        elif i % 9 == 2:
-            # 後手歩
-            p = 15
-        elif i == 72 or i == 0:
-            # 後手香車
-            p = 16
-        elif i == 63 or i == 9:
-            # 後手桂馬
-            p = 17
-        elif i == 54 or i == 18:
-            # 後手銀
-            p = 18
-        elif i == 10:
-            # 後手角
-            p = 19
-        elif i == 64:
-            # 後手飛車
-            p = 20
-        elif i == 45 or i == 27:
-            # 後手金
-            p = 21
-        elif i == 36:
-            # 後手玉
-            p = 22
-        else:
-            # 空白
-            p = 0
-        array = array.at[p, i].set(1)
-    return array
+    # 相手に合法手がない場合→詰み
+    is_mate = _is_mate(state)
+    state, reward, terminated = jax.lax.cond(
+        is_mate,
+        lambda: jax.lax.cond(
+            (_action.is_drop & ((_action.piece == 1) | (_action.piece == 15))),
+            # actionのis_dropがTrueかつpieceが歩の場合、打ち歩詰めで負け
+            lambda: (state, _turn_to_reward(state.turn), True),
+            # そうでなければ普通の詰みで勝ち
+            lambda: (state, _turn_to_reward(_another_color(state)), True),
+        ),
+        lambda: (state, 0, False),
+    )
+    # 反則負け
+    state, reward, terminated = jax.lax.cond(
+        (is_check | is_double_pawn | is_stuck),
+        lambda: (state2, _turn_to_reward(_another_color(state2)), True),
+        lambda: (state, reward, terminated),
+    )
+    # 反則負け
+    state, reward, terminated = jax.lax.cond(
+        (is_oob | is_illegal),
+        lambda: (state1, _turn_to_reward(_another_color(state1)), True),
+        lambda: (state, reward, terminated),
+    )
+    return state, reward, terminated
 
 
-# 盤面の情報をStateに変換
-def _make_board(bs: jnp.ndarray) -> ShogiState:
-    board = jnp.zeros((29, 81), dtype=jnp.int32)
-    for i in range(81):
-        board = board.at[0, i].set(0)
-        board = board.at[bs[i], i].set(1)
-    return ShogiState(board=board)  # type: ignore
+# turnから報酬計算
+def _turn_to_reward(turn: int):
+    return jnp.int32([1, -1])[turn]
 
 
 def _pawn_move(turn: int) -> jnp.ndarray:
@@ -199,12 +142,11 @@ def _silver_move(turn: int) -> jnp.ndarray:
 
 def _gold_move(turn: int) -> jnp.ndarray:
     array = jnp.zeros((9, 9), dtype=jnp.int32)
-    if turn == 0:
-        array = array.at[3, 3].set(1)
-        array = array.at[5, 3].set(1)
-    else:
-        array = array.at[3, 5].set(1)
-        array = array.at[5, 5].set(1)
+    array = jax.lax.cond(
+        turn == 0,
+        lambda: array.at[3, 3].set(1).at[5, 3].set(1),
+        lambda: array.at[3, 5].set(1).at[5, 5].set(1),
+    )
     array = array.at[4, 3].set(1)
     array = array.at[3, 4].set(1)
     array = array.at[4, 5].set(1)
@@ -242,31 +184,23 @@ def _point_to_location(point: int) -> Tuple[int, int]:
 
 # 端にいる駒の動けない地点へのフラグを折る
 def _cut_outside(array: jnp.ndarray, point: int) -> jnp.ndarray:
-    new_array = array
     u, d, l, r = _is_side(point)
     u2, d2 = _is_second_line(point)
     # (4, 4)での動きを基準にはみ出すところをカットする
-    if u:
-        new_array = new_array.at[:, 3].set(0)
-    if d:
-        new_array = new_array.at[:, 5].set(0)
-    if r:
-        new_array = new_array.at[3, :].set(0)
-    if l:
-        new_array = new_array.at[5, :].set(0)
-    if u2:
-        new_array = new_array.at[:, 2].set(0)
-    if d2:
-        new_array = new_array.at[:, 6].set(0)
-    return new_array
+    array = jax.lax.cond(u, lambda: array.at[:, 3].set(0), lambda: array)
+    array = jax.lax.cond(d, lambda: array.at[:, 5].set(0), lambda: array)
+    array = jax.lax.cond(r, lambda: array.at[3, :].set(0), lambda: array)
+    array = jax.lax.cond(l, lambda: array.at[5, :].set(0), lambda: array)
+    array = jax.lax.cond(u2, lambda: array.at[:, 2].set(0), lambda: array)
+    array = jax.lax.cond(d2, lambda: array.at[:, 6].set(0), lambda: array)
+    return array
 
 
 # 駒種と位置から到達できる地点をすべて生成
 def _action_board(array: jnp.ndarray, point: int) -> jnp.ndarray:
-    new_array = array
     y, t = _point_to_location(point)
-    new_array = _cut_outside(new_array, point)
-    return jnp.roll(new_array, (y - 4, t - 4), axis=(0, 1)).reshape(81)
+    array = _cut_outside(array, point)
+    return jnp.roll(array, (y - 4, t - 4), axis=(0, 1)).reshape(81)
 
 
 # 大駒・香車以外の位置ごとの移動できる地点を記録
@@ -318,89 +252,111 @@ def _is_same_declining(_from: int, to: int) -> bool:
 
 
 # from, からdirectionの方向に行ったときにtoに到達できるか
-def _is_same_line(from_: int, to: int, direction: int) -> bool:
+def _is_same_line(from_: int, to: int, direction: int):
     dir = direction % 10
-    if dir == 0 or dir == 5:
-        return _is_same_column(from_, to)
-    elif dir == 3 or dir == 4:
-        return _is_same_row(from_, to)
-    elif dir == 2 or dir == 6:
-        return _is_same_rising(from_, to)
-    elif dir == 1 or dir == 7:
-        return _is_same_declining(from_, to)
-    else:
-        return False
-
-
+    return jnp.bool_(
+        [
+            _is_same_column(from_, to),  # 0
+            _is_same_declining(from_, to),  # 1
+            _is_same_rising(from_, to),  # 2
+            _is_same_row(from_, to),  # 3
+            _is_same_row(from_, to),  # 4
+            _is_same_column(from_, to),  # 5
+            _is_same_rising(from_, to),  # 6
+            _is_same_declining(from_, to),  # 7
+            False,  # 8
+            False,  # 9
+        ]
+    )[dir]
+
+
+# TODO: cache
 # fromからある方向に移動させたときの位置と距離の関係
 # それ以上その方向に動かせないなら全部0
 def _dis_direction_array(from_: int, turn: int, direction: int) -> jnp.ndarray:
     array = jnp.zeros(81, dtype=jnp.int32)
     dif = _direction_to_dif(direction, turn)
-    for i in range(8):
+
+    def update_arr(i, array):
         dis = 1 + i
         to = from_ + dif * dis
-        if _is_in_board(to) and _is_same_line(from_, to, direction):
-            array = array.at[to].set(dis)
+        return jax.lax.cond(
+            _is_in_board(to) & _is_same_line(from_, to, direction),
+            lambda: array.at[to].set(dis),
+            lambda: array,
+        )
+
+    array = jax.lax.fori_loop(0, 8, update_arr, array)
     return array
 
 
+# TODO: cache
 # fromの座標とtoの座標からdirを生成
 def _point_to_direction(_from: int, to: int, promote: bool, turn: int) -> int:
-    direction = -1
-    dis = to - _from
     # 後手番の動きは反転させる
-    if turn == 1:
-        dis = -dis
+    dis = (to - _from) * jnp.int32([1, -1])[turn]
     # UP, UP_LEFT, UP_RIGHT, LEFT, RIGHT, DOWN, DOWN_LEFT, DOWN_RIGHT, UP2_LEFT, UP2_RIGHT, UP_PROMOTE...
     # の順でdirを割り振る
     # PROMOTEの場合は+10する処理を入れる
-    if _is_same_column(_from, to) and dis < 0:
-        direction = 0
-    if _is_same_declining(_from, to) and dis > 0:
-        direction = 1
-    if _is_same_rising(_from, to) and dis < 0:
-        direction = 2
-    if _is_same_row(_from, to) and dis > 0:
-        direction = 3
-    if _is_same_row(_from, to) and dis < 0:
-        direction = 4
-    if _is_same_column(_from, to) and dis > 0:
-        direction = 5
-    if _is_same_rising(_from, to) and dis > 0:
-        direction = 6
-    if _is_same_declining(_from, to) and dis < 0:
-        direction = 7
-    if dis == 7 and not _is_same_column(_from, to):
-        direction = 8
-    if dis == -11 and not _is_same_column(_from, to):
-        direction = 9
-    if promote:
-        direction += 10
+    direction = -1
+    direction = jax.lax.cond(
+        _is_same_column(_from, to) & (dis < 0), lambda: 0, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_declining(_from, to) & (dis > 0), lambda: 1, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_rising(_from, to) & (dis < 0), lambda: 2, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_row(_from, to) & (dis > 0), lambda: 3, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_row(_from, to) & (dis < 0), lambda: 4, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_column(_from, to) & (dis > 0), lambda: 5, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_rising(_from, to) & (dis > 0), lambda: 6, lambda: direction
+    )
+    direction = jax.lax.cond(
+        _is_same_declining(_from, to) & (dis < 0), lambda: 7, lambda: direction
+    )
+    direction = jax.lax.cond(
+        (dis == 7) & ~_is_same_column(_from, to), lambda: 8, lambda: direction
+    )
+    direction = jax.lax.cond(
+        (dis == -11) & ~_is_same_column(_from, to),
+        lambda: 9,
+        lambda: direction,
+    )
+    direction = jax.lax.cond(
+        promote, lambda: direction + 10, lambda: direction
+    )
     return direction
 
 
+# TODO: cache
 # 打った駒の種類をdirに変換
 def _hand_to_direction(piece: int) -> int:
     # 移動のdirはPROMOTE_UP_RIGHT2の19が最大なので20以降に配置
     # 20: 先手歩 21: 先手香車... 33: 後手金　に対応させる
-    if piece <= 14:
-        return 19 + piece
-    else:
-        return 12 + piece
+    return jax.lax.cond(piece <= 14, lambda: 19 + piece, lambda: 12 + piece)
 
 
 # ShogiActionをdlshogiのint型actionに変換
-def _action_to_dlaction(action: ShogiAction, turn: int) -> int:
-    if action.is_drop:
-        return _dlshogi_action(_hand_to_direction(action.piece), action.to)
-    else:
-        return _dlshogi_action(
+def _action_to_dlaction(action: ShogiAction, turn: int):
+    return jax.lax.cond(
+        action.is_drop,
+        lambda: _dlshogi_action(_hand_to_direction(action.piece), action.to),
+        lambda: _dlshogi_action(
             _point_to_direction(
                 action.from_, action.to, action.is_promote, turn
             ),
             action.to,
-        )
+        ),
+    )
 
 
 # dlshogiのint型actionをdirectionとtoに分解
@@ -410,32 +366,10 @@ def _separate_dlaction(action: int) -> Tuple[int, int]:
 
 
 # directionからfromがtoからどれだけ離れてるかを返す
-def _direction_to_dif(direction: int, turn: int) -> int:
-    dif = 0
-    if direction % 10 == 0:
-        dif = -1
-    if direction % 10 == 1:
-        dif = 8
-    if direction % 10 == 2:
-        dif = -10
-    if direction % 10 == 3:
-        dif = 9
-    if direction % 10 == 4:
-        dif = -9
-    if direction % 10 == 5:
-        dif = 1
-    if direction % 10 == 6:
-        dif = 10
-    if direction % 10 == 7:
-        dif = -8
-    if direction % 10 == 8:
-        dif = 7
-    if direction % 10 == 9:
-        dif = -11
-    if turn == 0:
-        return dif
-    else:
-        return -dif
+def _direction_to_dif(direction: int, turn: int):
+    direction = direction % 10
+    dif = jnp.int32([-1, 8, -10, 9, -9, 1, 10, -8, 7, -11])[direction]
+    return dif * jnp.int32([1, -1])[turn]
 
 
 # directionとto,stateから大駒含めた移動のfromの位置を割り出す
@@ -448,22 +382,25 @@ def _direction_to_from(
     _from = -1
     for i in range(8):
         f -= dif
-        if _is_in_board(f) and _from == -1 and _piece_type(state, f) != 0:
-            _from = f
+        _from = jax.lax.cond(
+            _is_in_board(f) & (_from == -1) & (_piece_type(state, f) != 0),
+            lambda: f,
+            lambda: _from,
+        )
     return _from, direction >= 10
 
 
+# TODO: cache
 def _direction_to_hand(direction: int) -> int:
-    if direction <= 26:
-        # direction:20が先手の歩（pieceの1）に対応
-        return direction - 19
-    else:
-        # direction:27が後手の歩（pieceの15）に対応
-        return direction - 12
+    return jax.lax.cond(
+        direction <= 26,
+        lambda: direction - 19,  # direction:20が先手の歩（pieceの1）に対応
+        lambda: direction - 12,  # direction:27が後手の歩（pieceの15）に対応
+    )
 
 
-def _piece_type(state: ShogiState, point: int) -> int:
-    return int(state.board[:, point].argmax())
+def _piece_type(state: ShogiState, point: int):
+    return state.board[:, point].argmax()
 
 
 # dlshogiのactionの情報をShogiActionの情報に変換
@@ -485,12 +422,11 @@ def _dlshogi_drop_action(
 
 def _dlaction_to_action(action: int, state: ShogiState) -> ShogiAction:
     direction, to = _separate_dlaction(action)
-    if direction <= 19:
-        # 移動のaction
-        return _dlshogi_move_action(direction, to, state)
-    else:
-        # 駒打ち
-        return _dlshogi_drop_action(direction, to, state)
+    return jax.lax.cond(
+        direction <= 19,
+        lambda: _dlshogi_move_action(direction, to, state),
+        lambda: _dlshogi_drop_action(direction, to, state),
+    )
 
 
 # 手番側でない色を返す
@@ -500,39 +436,31 @@ def _another_color(state: ShogiState) -> int:
 
 # pointが盤面内かどうか
 def _is_in_board(point: int) -> bool:
-    return 0 <= point <= 80
+    return (0 <= point) & (point <= 80)
 
 
 # 相手の駒を同じ種類の自分の駒に変換する
 def _convert_piece(piece: int) -> int:
-    if piece == 0:
-        return 0
     p = (piece + 14) % 28
-    if p == 0:
-        return 28
-    else:
-        return p
+    return jax.lax.cond(
+        piece == 0,
+        lambda: 0,
+        lambda: jax.lax.cond(p == 0, lambda: 28, lambda: p),
+    )
 
 
 # 駒から持ち駒への変換
 # 先手歩が0、後手金が13
-def _piece_to_hand(piece: int) -> int:
-    if piece % 14 == 0 or piece % 14 >= 9:
-        p = piece - 8
-    else:
-        p = piece
-    if p < 15:
-        return p - 1
-    else:
-        return p - 8
+def _piece_to_hand(piece: int):
+    p = jax.lax.cond(
+        (piece % 14 == 0) | (piece % 14 >= 9), lambda: piece - 8, lambda: piece
+    )
+    return jax.lax.cond(p < 15, lambda: p - 1, lambda: p - 8)
 
 
 # ある駒の持ち主を返す
 def _owner(piece):
-    if piece == 0:
-        return 2
-    else:
-        return (piece - 1) // 14
+    return jax.lax.cond(piece == 0, lambda: 2, lambda: (piece - 1) // 14)
 
 
 # 盤面のどこに何の駒があるかをnp.arrayに移したもの
@@ -543,11 +471,8 @@ def _board_status(state: ShogiState) -> jnp.ndarray:
 
 # 駒の持ち主の判定
 def _pieces_owner(state: ShogiState) -> jnp.ndarray:
-    board = jnp.zeros(81, dtype=jnp.int32)
-    for i in range(81):
-        piece = _piece_type(state, i)
-        board = board.at[i].set(_owner(piece))
-    return board
+    pieces = state.board.argmax(axis=0)
+    return jax.vmap(_owner)(pieces)
 
 
 # 駒の移動の盤面変換
@@ -560,12 +485,18 @@ def _move(
     board = board.at[action.piece, action.from_].set(0)
     board = board.at[0, action.from_].set(1)
     board = board.at[action.captured, action.to].set(0)
-    if action.is_promote:
-        board = board.at[action.piece + 8, action.to].set(1)
-    else:
-        board = board.at[action.piece, action.to].set(1)
-    if action.captured != 0:
-        hand = hand.at[_piece_to_hand(_convert_piece(action.captured))].add(1)
+    board = jax.lax.cond(
+        action.is_promote,
+        lambda: board.at[action.piece + 8, action.to].set(1),
+        lambda: board.at[action.piece, action.to].set(1),
+    )
+    hand = jax.lax.cond(
+        action.captured != 0,
+        lambda: hand.at[_piece_to_hand(_convert_piece(action.captured))].add(
+            1
+        ),
+        lambda: hand,
+    )
     return state.replace(board=board, hand=hand)  # type: ignore
 
 
@@ -581,26 +512,29 @@ def _drop(state: ShogiState, action: ShogiAction) -> ShogiState:
 def _inner_point(bs_one: jnp.ndarray, from_: int, direction: int):
     dir_array = _dis_direction_array(from_, 0, direction)
     # 途中で駒にぶつからない場合
-    if jnp.all(dir_array * bs_one == 0):
-        return from_ + _direction_to_dif(direction, 0) * jnp.max(dir_array)
-    else:
-        return _nearest_position(from_, direction, bs_one)
+    return jax.lax.cond(
+        jnp.all(dir_array * bs_one == 0),
+        lambda: from_ + _direction_to_dif(direction, 0) * jnp.max(dir_array),
+        lambda: _nearest_position(from_, direction, bs_one),
+    )
 
 
+# TODO: cache
 # fromからtoまでの地点をdifごとに1に置き換える
 # 最大8回
 # fromは0、toは1に
 def _change_between(from_: int, to: int, dif: int) -> jnp.ndarray:
     array = jnp.zeros(81, dtype=jnp.int32)
-    point = from_
-    flag = from_ != to
-    for i in range(8):
-        point += dif
-        if flag and _is_in_board(point):
-            array = array.at[point].set(1)
-        if point == to:
-            flag = False
-    return array
+    ix = jnp.arange(81)
+    return jax.lax.cond(
+        0 < dif,
+        lambda: jnp.where(
+            (from_ < ix) & (ix <= to) & ((ix - from_) % dif == 0), 1, array
+        ),
+        lambda: jnp.where(
+            (ix < from_) & (to <= ix) & ((ix - from_) % dif == 0), 1, array
+        ),
+    )
 
 
 # 香車の動き
@@ -659,69 +593,79 @@ def _dragon_move(bs: jnp.ndarray, from_: int) -> jnp.ndarray:
 def _piece_moves(bs: jnp.ndarray, piece, point) -> jnp.ndarray:
     moves = POINT_MOVES[point][piece]
     # 香車の動き
-    if piece == 2:
-        moves = _lance_move(bs, point, 0)
-    if piece == 16:
-        moves = _lance_move(bs, point, 1)
+    moves = jax.lax.cond(
+        piece == 2, lambda: _lance_move(bs, point, 0), lambda: moves
+    )
+    moves = jax.lax.cond(
+        piece == 16, lambda: _lance_move(bs, point, 1), lambda: moves
+    )
     # 角の動き
-    if piece == 5 or piece == 19:
-        moves = _bishop_move(bs, point)
+    is_bishop = (piece == 5) | (piece == 19)
+    moves = jax.lax.cond(
+        is_bishop, lambda: _bishop_move(bs, point), lambda: moves
+    )
     # 飛車の動き
-    if piece == 6 or piece == 20:
-        moves = _rook_move(bs, point)
+    is_rook = (piece == 6) | (piece == 20)
+    moves = jax.lax.cond(is_rook, lambda: _rook_move(bs, point), lambda: moves)
     # 馬の動き
-    if piece == 13 or piece == 27:
-        moves = _horse_move(bs, point)
+    is_horse = (piece == 13) | (piece == 27)
+    moves = jax.lax.cond(
+        is_horse, lambda: _horse_move(bs, point), lambda: moves
+    )
     # 龍の動き
-    if piece == 14 or piece == 28:
-        moves = _dragon_move(bs, point)
+    is_dragon = (piece == 14) | (piece == 28)
+    moves = jax.lax.cond(
+        is_dragon, lambda: _dragon_move(bs, point), lambda: moves
+    )
     return moves
 
 
 # 小駒のactionのみを返すpiece_moves
-def _small_piece_moves(piece: int, point: int) -> jnp.ndarray:
+def _small_piece_moves(piece, point) -> jnp.ndarray:
     return POINT_MOVES[point][piece]
 
 
 # 敵陣かどうか
-def _is_enemy_zone(turn: int, point: int) -> bool:
-    if turn == 0:
-        return point % 9 <= 2
-    else:
-        return point % 9 >= 6
+def _is_enemy_zone(turn: int, point: int):
+    return jnp.bool_([point % 9 <= 2, point % 9 >= 6])[turn]
 
 
 # 成れるかどうか
-def _can_promote(piece: int, _from: int, to: int) -> bool:
-    # pieceが飛車以下でないと成れない
-    if piece % 14 > 6 or piece % 14 == 0:
-        return False
-    # _fromとtoのどちらかが敵陣であれば成れる
-    else:
-        return _is_enemy_zone(_owner(piece), _from) or _is_enemy_zone(
-            _owner(piece), to
-        )
+def _can_promote(piece: int, _from: int, to: int):
+    can_promote = (piece % 14 <= 6) & (piece % 14 != 0)
+    can_promote &= _is_enemy_zone(_owner(piece), _from) | _is_enemy_zone(
+        _owner(piece), to
+    )
+    return can_promote
 
 
 def _create_one_piece_actions(
     piece: int, _from: int, to: int, actions: jnp.ndarray
 ) -> jnp.ndarray:
-    new_actions = actions
     normal_dir = _point_to_direction(_from, to, False, _owner(piece))
     normal_act = _dlshogi_action(normal_dir, to)
-    new_actions = new_actions.at[normal_act].set(1)
+    actions = actions.at[normal_act].set(1)
     pro_dir = _point_to_direction(_from, to, True, _owner(piece))
     pro_act = _dlshogi_action(pro_dir, to)
-    if _can_promote(piece, _from, to):
-        new_actions = new_actions.at[pro_act].set(1)
-    return new_actions
+    actions = jax.lax.cond(
+        _can_promote(piece, _from, to),
+        lambda: actions.at[pro_act].set(1),
+        lambda: actions,
+    )
+    return actions
 
 
 def _create_actions(piece, _from, to: jnp.ndarray) -> jnp.ndarray:
     actions = jnp.zeros(2754, dtype=jnp.int32)
-    for i in range(81):
-        if to[i] != 0:
-            actions = _create_one_piece_actions(piece, _from, i, actions)
+
+    def add_action(i, actions):
+        return jax.lax.cond(
+            to[i] != 0,
+            lambda: _create_one_piece_actions(piece, _from, i, actions),
+            lambda: actions,
+        )
+
+    actions = jax.lax.fori_loop(0, 81, add_action, actions)
     return actions
 
 
@@ -750,342 +694,377 @@ def _add_move_actions(piece, _from, array: jnp.ndarray) -> jnp.ndarray:
 
 
 # 駒の種類と位置から生成できるactionのフラグを折る
-def _filter_move_actions(
-    piece: int, _from: int, array: jnp.ndarray
-) -> jnp.ndarray:
+def _filter_move_actions(piece, _from, array: jnp.ndarray) -> jnp.ndarray:
     actions = _create_piece_actions(piece, _from)
     return jnp.where(actions == 1, 0, array)
 
 
 # 駒打ちのactionを追加する
 def _add_drop_actions(piece: int, array: jnp.ndarray) -> jnp.ndarray:
-    new_array = array
     direction = _hand_to_direction(piece)
-    new_array = new_array.at[
-        jnp.arange(81 * direction, 81 * (direction + 1))
-    ].set(1)
-    return new_array
+    idx = jnp.arange(ACTION_SIZE)
+    array = jnp.where(
+        (81 * direction <= idx) & (idx < 81 * (direction + 1)), 1, array
+    )
+    return array
 
 
 # 駒打ちのactionのフラグを折る
 def _filter_drop_actions(piece: int, array: jnp.ndarray) -> jnp.ndarray:
-    new_array = array
     direction = _hand_to_direction(piece)
-    new_array = new_array.at[
-        jnp.arange(81 * direction, 81 * (direction + 1))
-    ].set(0)
-    return new_array
+    ix = jnp.arange(2754)
+    array = jnp.where(
+        (81 * direction <= ix) & (ix < 81 * (direction + 1)), 0, array
+    )
+    return array
 
 
 # stateからblack,white両方のlegal_actionsを生成する
 # 普段は使わないがlegal_actionsが設定されていない場合に使用
 def _init_legal_actions(state: ShogiState) -> ShogiState:
-    bs = _board_status(state)
-    legal_actions_black, legal_actions_white = (
-        state.legal_actions_black,
-        state.legal_actions_white,
-    )
+    pieces = _board_status(state)
+
     # 移動の追加
-    for i in range(81):
-        piece = bs[i]
-        if piece <= 14:
-            legal_actions_black = _add_move_actions(
-                piece, i, legal_actions_black
-            )
-        else:
-            legal_actions_white = _add_move_actions(
-                piece, i, legal_actions_white
-            )
+    def add_black_moves(i, legal_actions_black):
+        return jax.lax.cond(
+            pieces[i] <= 14,
+            lambda: _add_move_actions(pieces[i], i, legal_actions_black),
+            lambda: legal_actions_black,
+        )
+
+    def add_white_moves(i, legal_actions_white):
+        return jax.lax.cond(
+            pieces[i] > 14,
+            lambda: _add_move_actions(pieces[i], i, legal_actions_white),
+            lambda: legal_actions_white,
+        )
+
+    legal_actions_black = jax.lax.fori_loop(
+        0, 81, add_black_moves, state.legal_actions_black
+    )
+    legal_actions_white = jax.lax.fori_loop(
+        0, 81, add_white_moves, state.legal_actions_white
+    )
+
+    def add_black_drops(i, legal_actions_black):
+        return jax.lax.cond(
+            state.hand[i] != 0,
+            lambda: _add_drop_actions(1 + i, legal_actions_black),
+            lambda: legal_actions_black,
+        )
+
+    def add_white_drops(i, legal_actions_white):
+        return jax.lax.cond(
+            state.hand[i + 7] != 0,
+            lambda: _add_drop_actions(15 + i, legal_actions_white),
+            lambda: legal_actions_white,
+        )
+
     # 駒打ちの追加
-    for i in range(7):
-        if state.hand[i] != 0:
-            legal_actions_black = _add_drop_actions(1 + i, legal_actions_black)
-        if state.hand[i + 7] != 0:
-            legal_actions_white = _add_drop_actions(
-                15 + i, legal_actions_white
-            )
+    legal_actions_black = jax.lax.fori_loop(
+        0, 7, add_black_drops, legal_actions_black
+    )
+    legal_actions_white = jax.lax.fori_loop(
+        0, 7, add_white_drops, legal_actions_white
+    )
     return state.replace(legal_actions_black=legal_actions_black, legal_actions_white=legal_actions_white)  # type: ignore
 
 
 # 成駒を成る前の駒に変更
 def _degeneration_piece(piece: int) -> int:
-    if piece % 14 >= 9 or piece == 14 or piece == 28:
-        return piece - 8
-    else:
-        return piece
+    return jax.lax.cond(
+        (piece % 14 >= 9) | (piece == 14) | (piece == 28),
+        lambda: piece - 8,
+        lambda: piece,
+    )
 
 
 # 駒の移動によるlegal_actionsの更新
 def _update_legal_move_actions(
     state: ShogiState, action: ShogiAction
 ) -> ShogiState:
-    s = state
-    if s.turn == 0:
-        player_actions = s.legal_actions_black
-        enemy_actions = s.legal_actions_white
-    else:
-        player_actions = s.legal_actions_white
-        enemy_actions = s.legal_actions_black
+    player_actions, enemy_actions = jax.lax.cond(
+        state.turn == 0,
+        lambda: (state.legal_actions_black, state.legal_actions_white),
+        lambda: (state.legal_actions_white, state.legal_actions_black),
+    )
     # 元の位置にいたときのフラグを折る
-    new_player_actions = _filter_move_actions(
+    player_actions = _filter_move_actions(
         action.piece, action.from_, player_actions
     )
-    new_enemy_actions = enemy_actions
     # 移動後の位置からの移動のフラグを立てる
-    new_piece = action.piece
-    if action.is_promote:
-        new_piece += 8
-    new_player_actions = _add_move_actions(
-        new_piece, action.to, new_player_actions
+    piece = jax.lax.cond(
+        action.is_promote, lambda: action.piece + 8, lambda: action.piece
     )
+    player_actions = _add_move_actions(piece, action.to, player_actions)
     # 取った駒を自分の持ち駒に変換
     # 取っていない場合は0
     captured = _degeneration_piece(_convert_piece(action.captured))
     # 駒が取られた場合、相手の取られた駒によってできていたactionのフラグを折る
-    if action.captured != 0:
-        new_enemy_actions = _filter_move_actions(
-            action.captured, action.to, new_enemy_actions
-        )
-        new_player_actions = _add_drop_actions(captured, new_player_actions)
-    if s.turn == 0:
-        legal_actions_black, legal_actions_white = (
-            new_player_actions,
-            new_enemy_actions,
-        )
-    else:
-        legal_actions_black, legal_actions_white = (
-            new_enemy_actions,
-            new_player_actions,
-        )
-    return s.replace(legal_actions_black=legal_actions_black, legal_actions_white=legal_actions_white)  # type: ignore
+    enemy_actions = jax.lax.cond(
+        action.captured != 0,
+        lambda: _filter_move_actions(
+            action.captured, action.to, enemy_actions
+        ),
+        lambda: enemy_actions,
+    )
+    player_actions = jax.lax.cond(
+        action.captured != 0,
+        lambda: _add_drop_actions(captured, player_actions),
+        lambda: player_actions,
+    )
+
+    legal_actions_black, legal_actions_white = jax.lax.cond(
+        state.turn == 0,
+        lambda: (player_actions, enemy_actions),
+        lambda: (enemy_actions, player_actions),
+    )
+    return state.replace(legal_actions_black=legal_actions_black, legal_actions_white=legal_actions_white)  # type: ignore
 
 
 # 駒打ちによるlegal_actionsの更新
 def _update_legal_drop_actions(
     state: ShogiState, action: ShogiAction
 ) -> ShogiState:
-    s = state
-    if s.turn == 0:
-        player_actions = s.legal_actions_black
-    else:
-        player_actions = s.legal_actions_white
-    # 移動後の位置からの移動のフラグを立てる
-    new_player_actions = _add_move_actions(
-        action.piece, action.to, player_actions
+    player_actions = jax.lax.cond(
+        state.turn == 0,
+        lambda: state.legal_actions_black,
+        lambda: state.legal_actions_white,
     )
+    # 移動後の位置からの移動のフラグを立てる
+    player_actions = _add_move_actions(action.piece, action.to, player_actions)
     # 持ち駒がもうない場合、その駒を打つフラグを折る
-    if s.hand[_piece_to_hand(action.piece)] == 1:
-        new_player_actions = _filter_drop_actions(
-            action.piece, new_player_actions
-        )
-    if s.turn == 0:
-        s = s.replace(legal_actions_black=new_player_actions)  # type: ignore
-    else:
-        s = s.replace(legal_actions_white=new_player_actions)  # type: ignore
-    return s
+    player_actions = jax.lax.cond(
+        state.hand[_piece_to_hand(action.piece)] == 1,
+        lambda: _filter_drop_actions(action.piece, player_actions),
+        lambda: player_actions,
+    )
+    return jax.lax.cond(
+        state.turn == 0,
+        lambda: state.replace(legal_actions_black=player_actions),  # type: ignore
+        lambda: state.replace(legal_actions_white=player_actions),  # type: ignore
+    )
 
 
 # 自分の駒がある位置への移動を除く
 def _filter_my_piece_move_actions(
     turn: int, owner: jnp.ndarray, array: jnp.ndarray
 ) -> jnp.ndarray:
-    new_array = array
-    for i in range(81):
-        if owner[i] != turn:
-            continue
-        new_array = new_array.at[jnp.arange(i, 1620 + i, 81)].set(0)
-    return new_array
+    actions = array.reshape(34, 81)
+    mask = jnp.tile(owner == turn, reps=(34, 1))
+    mask = mask.at[20:].set(False)
+    actions = jnp.where(mask, 0, actions)
+    return actions.flatten()
 
 
 # 駒がある地点への駒打ちを除く
 def _filter_occupied_drop_actions(
     turn: int, owner: jnp.ndarray, array: jnp.ndarray
 ) -> jnp.ndarray:
-    new_array = array
-    for i in range(81):
-        if owner[i] == 2:
-            continue
-        if turn == 0:
-            new_array = new_array.at[
-                jnp.arange(81 * 20 + i, 81 * 27 + i, 81)
-            ].set(0)
-        else:
-            new_array = new_array.at[
-                jnp.arange(81 * 27 + i, 81 * 34 + i, 81)
-            ].set(0)
-    return new_array
+    """
+    x x x x x x x x x
+    x x x x x x x x x
+    --- 20
+    x x x x F x x x x
+    x x x x F x x x x
+    --- 27
+    x x x x x x x x x
+    x x x x x x x x x
+    """
+    actions = array.reshape(34, 81)
+    mask1 = jnp.tile(owner != 2, reps=(34, 1))
+    s = jnp.int32([20, 27])[turn]
+    ix = jnp.arange(34)
+    mask2 = jnp.tile((s <= ix) & (ix < s + 7), (81, 1)).transpose()
+    mask = mask1 & mask2
+    actions = jnp.where(mask, 0, actions)
+    return actions.flatten()
 
 
 # boardのlegal_actionsを利用して合法手を生成する
 # 大駒や香車の利きはboardのlegal_actionsに追加していないので、ここで追加する
 # 自殺手や反則手はここでは除かない
 def _legal_actions(state: ShogiState) -> jnp.ndarray:
-    if state.turn == 0:
-        action_array = state.legal_actions_black
-    else:
-        action_array = state.legal_actions_white
-    bs = _board_status(state)
+    actions = jax.lax.cond(
+        state.turn == 0,
+        lambda: state.legal_actions_black,
+        lambda: state.legal_actions_white,
+    )
+    pieces = _board_status(state)
     own = _pieces_owner(state)
-    for i in range(81):
-        piece = bs[i]
-        if _owner(piece) == state.turn:
-            action_array = _add_action(
-                _create_actions(piece, i, _piece_moves(bs, piece, i)),
-                action_array,
-            )
+
+    # TODO: vectorize
+    def add_action(i, actions):
+        return jax.lax.cond(
+            _owner(pieces[i]) == state.turn,
+            lambda: _add_action(
+                _create_actions(
+                    pieces[i], i, _piece_moves(pieces, pieces[i], i)
+                ),
+                actions,
+            ),
+            lambda: actions,
+        )
+
+    actions = jax.lax.fori_loop(0, 81, add_action, actions)
     # 自分の駒がある位置への移動actionを除く
-    action_array = _filter_my_piece_move_actions(state.turn, own, action_array)
+    actions = _filter_my_piece_move_actions(state.turn, own, actions)
     # 駒がある地点への駒打ちactionを除く
-    action_array = _filter_occupied_drop_actions(state.turn, own, action_array)
-    return action_array
+    actions = _filter_occupied_drop_actions(state.turn, own, actions)
+    return actions
 
 
 # 王手判定
 # 密接・遠隔の王手で分ける
 def _is_check(state: ShogiState) -> Tuple[int, jnp.ndarray, int, jnp.ndarray]:
-    # そもそも王がいない場合
-    if jnp.all(state.board[8 + 14 * state.turn] == 0):
-        return (
+    return jax.lax.cond(
+        jnp.all(state.board[8 + 14 * state.turn] == 0),
+        # そもそも王がいない場合
+        lambda: (
             0,
             jnp.zeros(81, dtype=jnp.int32),
             0,
             jnp.zeros(81, dtype=jnp.int32),
-        )
-    else:
-        return _is_check_(state)
+        ),
+        lambda: _is_check_(state),
+    )
 
 
 # 玉がいる前提
-def _is_check_(state: ShogiState) -> Tuple[int, jnp.ndarray, int, jnp.ndarray]:
+def _is_check_(state: ShogiState):
     check = 0
     checking_point = jnp.zeros((2, 81), dtype=jnp.int32)
-    king_point = int(state.board[8 + 14 * state.turn, :].argmax())
+    king_point = state.board[8 + 14 * state.turn, :].argmax()
     near_king = _small_piece_moves(8 + 14 * state.turn, king_point)
-    bs = _board_status(state)
-    for i in range(81):
-        piece = bs[i]
-        if (
-            _owner(piece) == _another_color(state)
-            and _piece_moves(bs, piece, i)[king_point] == 1
-        ):
-            # 桂馬の王手も密接としてカウント
-            if near_king[i] == 1 or piece % 14 == 3:
-                check += 10
-                checking_point = checking_point.at[0, i].set(1)
-            else:
-                # 遠隔の王手は9以上ありえない
-                check += 1
-                checking_point = checking_point.at[1, i].set(1)
+    pieces = _board_status(state)
+
+    def update_check(i, x):
+        check, checking_point = x
+        flag = (_owner(pieces[i]) == _another_color(state)) & (
+            _piece_moves(pieces, pieces[i], i)[king_point] == 1
+        )
+        # 桂馬の王手も密接としてカウント
+        flag2 = (near_king[i] == 1) | (pieces[i] % 14 == 3)
+        check = jax.lax.cond(
+            flag,
+            lambda: jax.lax.cond(
+                flag2, lambda: check + 10, lambda: check + 1  # 遠隔の王手は9以上ありえない
+            ),
+            lambda: check,
+        )
+        checking_point = jax.lax.cond(
+            flag,
+            lambda: jax.lax.cond(
+                flag2,
+                lambda: checking_point.at[0, i].set(1),
+                lambda: checking_point.at[1, i].set(1),  # 遠隔の王手は9以上ありえない
+            ),
+            lambda: checking_point,
+        )
+        return check, checking_point
+
+    check, checking_point = jax.lax.fori_loop(
+        0, 81, update_check, (check, checking_point)
+    )
+
     return check // 10, checking_point[0], check % 10, checking_point[1]
 
 
 # 二歩判定
 # 手番側でチェックする
-def _is_double_pawn(state: ShogiState) -> bool:
-    is_double_pawn = False
+def _is_double_pawn(state: ShogiState):
     bs = _board_status(state)
-    for i in range(9):
-        num_pawn = jnp.count_nonzero(
-            bs[9 * i : 9 * (i + 1)] == 1 + state.turn * 14
-        )
-        if num_pawn >= 2:
-            is_double_pawn = True
-    return is_double_pawn
+    is_pawn = bs.reshape(9, 9) == 1 + state.turn * 14
+    return ((is_pawn).sum(axis=1) >= 2).any()
 
 
 # 行き所のない駒判定
-def _is_stuck(state: ShogiState) -> bool:
-    is_stuck = False
+def _is_stuck(state: ShogiState):
+    is_stuck = jnp.bool_(False)
     bs = _board_status(state)
     line1 = bs[0::9]
-    if jnp.any(line1 == 1):
-        is_stuck = True
-    if jnp.any(line1 == 2):
-        is_stuck = True
-    if jnp.any(line1 == 3):
-        is_stuck = True
+    is_stuck |= jnp.all(jnp.any(line1 == 1))
+    is_stuck |= jnp.all(jnp.any(line1 == 2))
+    is_stuck |= jnp.all(jnp.any(line1 == 3))
     line2 = bs[1::9]
-    if jnp.any(line2 == 3):
-        is_stuck = True
+    is_stuck |= jnp.all(jnp.any(line2 == 3))
     line8 = bs[7::9]
-    if jnp.any(line8 == 17):
-        is_stuck = True
+    is_stuck |= jnp.all(jnp.any(line8 == 17))
     line9 = bs[8::9]
-    if jnp.any(line9 == 15):
-        is_stuck = True
-    if jnp.any(line9 == 16):
-        is_stuck = True
-    if jnp.any(line9 == 17):
-        is_stuck = True
+    is_stuck |= jnp.all(jnp.any(line9 == 15))
+    is_stuck |= jnp.all(jnp.any(line9 == 16))
+    is_stuck |= jnp.all(jnp.any(line9 == 17))
     return is_stuck
 
 
 # pinされている方向から縦、横、右斜めうえ、右斜め下の四方向に分類
-def _direction_to_pin(direction: int) -> int:
-    if direction == 0 or direction == 5:
-        return 1
-    if direction == 1 or direction == 7:
-        return 2
-    if direction == 2 or direction == 6:
-        return 3
-    if direction == 3 or direction == 4:
-        return 4
-    return 0
+def _direction_to_pin(direction: int):
+    # assert direction < 8
+    return jnp.int32([1, 2, 3, 4, 4, 1, 3, 2])[direction]
 
 
 # それぞれの方向について、1番fromに近い駒の位置を返す
-def _nearest_position(from_: int, direction: int, bs_one: jnp.ndarray) -> int:
-    flag1 = True
-    point1 = -1
+def _nearest_position(from_: int, direction: int, bs_one: jnp.ndarray):
     dif = _direction_to_dif(direction, 0)
-    for i in range(8):
-        point = from_ + (1 + i) * dif
-        if (
-            flag1
-            and _is_in_board(point)
-            and _is_same_line(from_, point, direction)
-            and bs_one[point] == 1
-        ):
-            point1 = point
-            flag1 = False
-    return point1
+
+    @jax.vmap
+    def is_ok(point):
+        return (
+            _is_in_board(point)
+            & _is_same_line(from_, point, direction)
+            & (bs_one[point] == 1)
+        )
+
+    idxs = jnp.arange(8)
+    points = from_ + (1 + idxs) * dif
+    is_oks = is_ok(points)
+
+    min_ix = jnp.argmax(is_oks)
+    return points[min_ix]
 
 
 # pinされている駒の位置と方向を記録(1方向のみ)
 def _direction_pin(
     bs: jnp.ndarray,
-    turn: int,
-    king_point: int,
-    direction: int,
+    turn,
+    king_point,
+    direction,
     array: jnp.ndarray,
 ) -> jnp.ndarray:
-    new_array = array
     e_turn = (turn + 1) % 2
     bs_one = jnp.where(bs == 0, 0, 1)
     # 玉に一番近い駒の位置
     point1 = _nearest_position(king_point, direction, bs_one)
     # 2番目に近い駒の位置
     point2 = _nearest_position(point1, direction, bs_one)
-    if point1 == -1 or point2 == -1:
-        return new_array
     piece1 = bs[point1]
     piece2 = bs[point2]
-    if _owner(piece1) != turn:
-        return new_array
-    flag = False
-    if piece2 == 2 + 14 * e_turn and direction == 5 * turn:
-        flag = True
-    if (piece2 == 6 + e_turn * 14 or piece2 == 14 + e_turn * 14) and (
-        direction == 0 or direction == 3 or direction == 4 or direction == 5
-    ):
-        flag = True
-    if (piece2 == 5 + e_turn * 14 or piece2 == 13 + e_turn * 14) and (
-        direction == 1 or direction == 2 or direction == 6 or direction == 7
-    ):
-        flag = True
-    if flag:
-        new_array = new_array.at[point1].set(_direction_to_pin(direction))
-    return new_array
+    arr1 = array
+
+    flag = (piece2 == 2 + 14 * e_turn) & (direction == 5 * turn)
+    flag |= ((piece2 == 6 + e_turn * 14) | (piece2 == 14 + e_turn * 14)) & (
+        (direction == 0)
+        | (direction == 3)
+        | (direction == 4)
+        | (direction == 5)
+    )
+    flag |= ((piece2 == 5 + e_turn * 14) | (piece2 == 13 + e_turn * 14)) & (
+        (direction == 1)
+        | (direction == 2)
+        | (direction == 6)
+        | (direction == 7)
+    )
+    arr2 = jax.lax.cond(
+        flag,
+        lambda: array.at[point1].set(_direction_to_pin(direction)),
+        lambda: array,
+    )
+
+    return jax.lax.cond(
+        ((point1 == -1) | (point2 == -1)) | (_owner(piece1) != turn),
+        lambda: arr1,
+        lambda: arr2,
+    )
 
 
 # pinされている駒の位置と方向を記録
@@ -1093,125 +1072,123 @@ def _pin(state: ShogiState) -> jnp.ndarray:
     bs = _board_status(state)
     turn = state.turn
     pins = jnp.zeros(81, dtype=jnp.int32)
-    king_point = int(state.board[8 + 14 * turn, :].argmax())
-    for i in range(8):
-        pins = _direction_pin(bs, turn, king_point, i, pins)
+    king_point = state.board[8 + 14 * turn, :].argmax()
+    pins = jax.lax.fori_loop(
+        0, 8, lambda i, p: _direction_pin(bs, turn, king_point, i, p), pins
+    )
     return pins
 
 
 # 特定の方向の動きだけを除いたactionを生成する
 def _eliminate_direction(actions: jnp.ndarray, direction) -> jnp.ndarray:
-    new_array = actions
+    actions = actions.reshape(34, 81)
     # 2方向と成・不成の4方向のフラグを折る
-    dir1 = 0
-    dir2 = 0
-    # 縦
-    if direction == 1:
-        dir1 = 0
-        dir2 = 5
-    # 右下がり
-    if direction == 2:
-        dir1 = 1
-        dir2 = 7
-    # 右上がり
-    if direction == 3:
-        dir1 = 2
-        dir2 = 6
-    # 横
-    if direction == 4:
-        dir1 = 3
-        dir2 = 4
+    dir1, dir2 = 0, 0
+    dir1, dir2 = jax.lax.switch(
+        direction - 1,
+        [
+            lambda: (0, 5),  # direction == 1, 縦
+            lambda: (1, 7),  # direction == 2, 右下がり
+            lambda: (2, 6),  # direction == 3, 右上り
+            lambda: (3, 4),  # direction == 4, 横
+        ],
+    )
     pro_dir1 = dir1 + 8
     pro_dir2 = dir2 + 8
-    new_array = new_array.at[dir1 * 81 : (dir1 + 1) * 81].set(0)
-    new_array = new_array.at[dir2 * 81 : (dir2 + 1) * 81].set(0)
-    new_array = new_array.at[pro_dir1 * 81 : (pro_dir1 + 1) * 81].set(0)
-    new_array = new_array.at[pro_dir2 * 81 : (pro_dir2 + 1) * 81].set(0)
-    return new_array
+    actions = actions.at[dir1].set(0)
+    actions = actions.at[dir2].set(0)
+    actions = actions.at[pro_dir1].set(0)
+    actions = actions.at[pro_dir2].set(0)
+    return actions.flatten()
 
 
 # 利きの判定
 # 玉の位置を透過する（玉をいないものとして扱う）ことで香車や角などの利きを玉の奥まで通す
 def _kingless_effected_positions(
-    bs: jnp.ndarray, king_point: int, turn: int
+    bs: jnp.ndarray, king_point, turn
 ) -> jnp.ndarray:
-    all_effect = jnp.zeros(81, dtype=jnp.int32)
     bs = bs.at[king_point].set(0)
-    for i in range(81):
-        if _owner(bs[i]) == turn:
-            all_effect += _piece_moves(bs, bs[i], i)
+    ix = jnp.arange(81)
+    owners = jax.vmap(_owner)(bs)
+    effects = jax.vmap(partial(_piece_moves, bs=bs))(piece=bs, point=ix)
+    mask = jnp.tile(owners == turn, reps=(81, 1)).transpose()
+    effects = jnp.where(mask, effects, 0)
+    all_effect = effects.sum(axis=0)
     return all_effect
 
 
 # 玉が移動する手に合法なものがあるかを調べる
-def _king_escape(state: ShogiState) -> bool:
-    king_point = int(state.board[8 + 14 * state.turn, :].argmax())
+def _king_escape(state: ShogiState):
+    king_point = state.board[8 + 14 * state.turn, :].argmax()
     bs = _board_status(state)
     effects = _kingless_effected_positions(
         bs, king_point, _another_color(state)
     )
     king_moves = POINT_MOVES[king_point][8]
-    flag = False
-    for i in range(81):
-        if (
-            king_moves[i] == 1
-            and _owner(bs[i]) != state.turn
-            and effects[i] == 0
-        ):
-            flag = True
-    return flag
+    owners = jax.vmap(_owner)(bs)
+    ok = ((king_moves == 1) & (owners != state.turn) & (effects == 0)).any(0)
+    return ok
 
 
 # point同士の間の位置にフラグを立てる
 def _between(point1: int, point2: int) -> jnp.ndarray:
     direction = _point_to_direction(point1, point2, False, 0)
-    if direction == -1:
-        bet = jnp.zeros(81, dtype=jnp.int32)
-    else:
-        bet = _change_between(point1, point2, _direction_to_dif(direction, 0))
+    bet = jax.lax.cond(
+        direction == -1,
+        lambda: jnp.zeros(81, dtype=jnp.int32),
+        lambda: _change_between(
+            point1, point2, _direction_to_dif(direction, 0)
+        ),
+    )
     bet = bet.at[point2].set(0)
     return bet
 
 
 # pinされている駒の非合法手を弾いてlegal_actionを返す
 def _eliminate_pin_actions(
-    bs: jnp.ndarray, pins: jnp.ndarray, l_actions: jnp.ndarray
+    bs: jnp.ndarray, pins: jnp.ndarray, actions: jnp.ndarray
 ):
-    for i in range(81):
-        if pins[i] != 0:
-            l_actions = _filter_action(
+    actions = jax.lax.fori_loop(
+        0,
+        81,
+        lambda i, a: jax.lax.cond(
+            pins[i] != 0,
+            lambda: _filter_action(
                 _eliminate_direction(_create_piece_actions(bs[i], i), pins[i]),
-                l_actions,
-            )
-    return l_actions
+                a,
+            ),
+            lambda: a,
+        ),
+        actions,
+    )
+    return actions
 
 
 # 玉が逃げる以外の手に王手回避の手があるかをチェック
 # 存在しない場合True
 # 両王手は考えない(事前にはじく)
 def _is_avoid_check(
-    cn: int,
-    cnp: jnp.ndarray,
-    cfp: jnp.ndarray,
-    king_point: int,
+    cn,
+    cnp,
+    cfp,
+    king_point,
     legal_actions: jnp.ndarray,
 ):
-    # 密接の王手
-    if cn == 1:
-        # 玉が逃げる手以外の合法手は王手をかけた駒がある座標への移動のみ
-        point = int(cnp.argmax())
-        points = cnp
-    # 開き王手
-    else:
-        point = int(cfp.argmax())
-        # pointとking_pointの間。ここに駒を打ったり移動させたりする手は合法
-        points = _between(king_point, point)
-    for i in range(81):
-        if points[i] == 0 and point != i:
-            legal_actions = legal_actions.at[
-                jnp.arange(i, 81 * 34 + i, 81)
-            ].set(0)
-    return (legal_actions == 0).all()
+    is_close_check = cn == 1  # 密接の王手 (or 開き王手）
+    # 玉が逃げる手以外の合法手は王手をかけた駒がある座標への移動のみ
+    point = jax.lax.cond(
+        is_close_check, lambda: cnp.argmax(), lambda: cfp.argmax()
+    )
+    # pointとking_pointの間。ここに駒を打ったり移動させたりする手は合法
+    points = jax.lax.cond(
+        is_close_check, lambda: cnp, lambda: _between(king_point, point)
+    )
+
+    actions = legal_actions.reshape(34, 81)
+    ix = jnp.arange(81)
+    mask = jnp.tile((points == 0) & (point != ix), reps=(34, 1))
+    actions = jnp.where(mask, 0, actions)
+    return (actions == 0).all()
 
 
 # 詰み判定関数
@@ -1219,10 +1196,11 @@ def _is_avoid_check(
 def _is_mate(state: ShogiState) -> bool:
     cn, cnp, cf, cfp = _is_check(state)
     # 王手がかかっていないならFalse
-    if cn + cf == 0:
-        return False
-    else:
-        return _is_mate_noncheck(cn, cnp, cf, cfp, state)
+    return jax.lax.cond(
+        cn + cf == 0,
+        lambda: FALSE,
+        lambda: _is_mate_noncheck(cn, cnp, cf, cfp, state),
+    )
 
 
 # 王手の有無にかかわらず詰みを判定する
@@ -1231,18 +1209,18 @@ def _is_mate_noncheck(
 ):
     legal_actions = _legal_actions(state)
     bs = _board_status(state)
-    king_point = int(state.board[8 + 14 * state.turn, :].argmax())
+    king_point = state.board[8 + 14 * state.turn, :].argmax()
     # 玉が逃げる手以外の合法手
     legal_actions = _filter_move_actions(
         8 + 14 * state.turn, king_point, legal_actions
     )
     # ピンされている駒の非合法な動きをのぞく
     legal_actions = _eliminate_pin_actions(bs, _pin(state), legal_actions)
-    # 玉が逃げる手が合法なら詰んでない
-    if _king_escape(state):
-        return False
-    else:
+    return jax.lax.cond(
+        _king_escape(state),
+        # 玉が逃げる手が合法なら詰んでない
+        lambda: FALSE,
         # 両王手がかかっている場合はTrue, そうでなければis_avoid_check参照
-        return cn + cf >= 2 or _is_avoid_check(
-            cn, cnp, cfp, king_point, legal_actions
-        )
+        lambda: (cn + cf >= 2)
+        | _is_avoid_check(cn, cnp, cfp, king_point, legal_actions),
+    )
diff --git a/workspace/shogi-cache.py b/workspace/shogi-cache.py
new file mode 100755
--- /dev/null
+++ b/workspace/shogi-cache.py
@@ -0,0 +1,66 @@
+#! python3
+import jax.numpy as jnp
+from flax.serialization import to_bytes
+
+
+# 初期盤面生成
+def _make_init_board() -> jnp.ndarray:
+    array = jnp.zeros((29, 81), dtype=jnp.int32)
+    for i in range(81):
+        if i % 9 == 6:
+            # 先手歩
+            p = 1
+        elif i == 8 or i == 80:
+            # 先手香車
+            p = 2
+        elif i == 17 or i == 71:
+            # 先手桂馬
+            p = 3
+        elif i == 26 or i == 62:
+            # 先手銀
+            p = 4
+        elif i == 70:
+            # 先手角
+            p = 5
+        elif i == 16:
+            # 先手飛車
+            p = 6
+        elif i == 35 or i == 53:
+            # 先手金
+            p = 7
+        elif i == 44:
+            # 先手玉
+            p = 8
+        elif i % 9 == 2:
+            # 後手歩
+            p = 15
+        elif i == 72 or i == 0:
+            # 後手香車
+            p = 16
+        elif i == 63 or i == 9:
+            # 後手桂馬
+            p = 17
+        elif i == 54 or i == 18:
+            # 後手銀
+            p = 18
+        elif i == 10:
+            # 後手角
+            p = 19
+        elif i == 64:
+            # 後手飛車
+            p = 20
+        elif i == 45 or i == 27:
+            # 後手金
+            p = 21
+        elif i == 36:
+            # 後手玉
+            p = 22
+        else:
+            # 空白
+            p = 0
+        array = array.at[p, i].set(1)
+    return array
+
+
+init_board = _make_init_board()
+print(f"INIT_BOARD = {to_bytes(init_board)}")


Test patch:
diff --git a/tests/test_shogi.py b/tests/test_shogi.py
--- a/tests/test_shogi.py
+++ b/tests/test_shogi.py
@@ -1,11 +1,41 @@
 from pgx.shogi import init, _action_to_dlaction, _dlaction_to_action, ShogiAction, ShogiState, _move, _drop, \
     _piece_moves, _is_check, _legal_actions, _add_drop_actions, _init_legal_actions, _update_legal_move_actions, \
-    _update_legal_drop_actions, _is_double_pawn, _is_stuck, _board_status, step, _between, _pin, _make_board, \
+    _update_legal_drop_actions, _is_double_pawn, _is_stuck, _board_status, step, _between, _pin, \
     _is_mate
 
 
+import jax
 import jax.numpy as jnp
 
+init = jax.jit(init)
+step = jax.jit(step)
+_action_to_dlaction = jax.jit(_action_to_dlaction)
+_dlaction_to_action = jax.jit(_dlaction_to_action)
+_move = jax.jit(_move)
+_drop = jax.jit(_drop)
+_piece_moves = jax.jit(_piece_moves)
+_is_check = jax.jit(_is_check)
+_legal_actions = jax.jit(_legal_actions)
+_add_drop_actions = jax.jit(_add_drop_actions)
+_init_legal_actions = jax.jit(_init_legal_actions)
+_update_legal_move_actions = jax.jit(_update_legal_move_actions)
+_update_legal_drop_actions = jax.jit(_update_legal_drop_actions)
+_is_double_pawn = jax.jit(_is_double_pawn)
+_is_stuck = jax.jit(_is_stuck)
+_board_status = jax.jit(_board_status)
+_between = jax.jit(_between)
+_pin = jax.jit(_pin)
+_is_mate = jax.jit(_is_mate)
+
+
+# 盤面の情報をStateに変換
+def _make_board(bs: jnp.ndarray) -> ShogiState:
+    board = jnp.zeros((29, 81), dtype=jnp.int32)
+    for i in range(81):
+        board = board.at[0, i].set(0)
+        board = board.at[bs[i], i].set(1)
+    return ShogiState(board=board)  # type: ignore
+
 
 def make_test_board():
     board = jnp.zeros((29, 81), dtype=jnp.int32)