`jax.jit` を使って高速化するときのメモ about pgx

[OkanoShinri]

forはコンパイル時に全て展開されてかなり時間を食うので、固定長であっても使わない方が良い。
jax.lax.fori_loopやjax.lax.mapに変える。

• before

# 遅い
for xy in range(BOARD_SIZE * BOARD_SIZE):
  _, _, is_illegal = step(state, xy)
  legal_action = legal_action.at[xy].set(is_illegal)

• after

# 速い
jax.lax.map(lambda xy: step(state, xy), jnp.arange(BOARD_SIZE * BOARD_SIZE))[2]

---

どうしてもforですべての要素を調べないといけない場合が生じる場合、アルゴリズムを見直した方が良い。
（stateに新たに変数を追加するくらいで遅くなったりはしない）
例：囲碁で石を取った時に、隣接する連に呼吸点を追加する部分

• before

for _xy in range(BOARD_SIZE * BOARD_SIZE):
  for _around in [[-1, 0], [1, 0], [0, 1], [0, -1]]:
    # 全ての点の周囲四方を調べ、取られた石と一致したら呼吸点に追加

• after

# 隣接しているかどうかをstateに記録(0:なし 1:呼吸点 2:石)
jax.lax.map(
  lambda l: jnp.where((l > 0) & surrounded_stones, 1, l),  # 呼吸点に追加
  _state.liberty,
)

---

whileは条件が合えば使える。
jaxの制限を意識しすぎて非効率的なアルゴリズムを組むくらいならjax.lax.while_loopを使った方が良い。

[sotetsuk]

関連がありそうなリンク: https://medium.com/@ngoodger_7766/writing-an-rl-environment-in-jax-9f74338898ba

[sotetsuk]

関数の途中でのリターンやbreakはできる？どうする？

[sotetsuk]

• 途中でのリターンはcondで分けてreturn以降を別関数へ分ける
• fori_loop, map, scanを使い分ける
• while_loop
• どうすてもifがネストするときには条件に名前を付けてネストを浅くする