[Backgammon] legal_action_mask: scatter once instead of a (26, 156) intermediate

[gweber]

What

_legal_action_mask_for_valid_single_dice computed the per-die legal mask like this:

def _is_legal(idx):
    action = idx * 6 + die
    m = jnp.zeros(26 * 6, dtype=jnp.bool_)      # a full 156-wide zero vector ...
    return m.at[action].set(_is_action_legal(board, action))   # ... to set a single bit
legal = jax.vmap(_is_legal)(jnp.arange(26)).any(axis=0)

i.e. it materialized a (26, 26*6) intermediate (one zero vector per src, each with a single bit set) and OR-reduced it. This computes the 26 src legalities as a (26,) vector and scatters them into the mask in one shot:

actions = jnp.arange(26, dtype=jnp.int32) * 6 + die
legal = jax.vmap(lambda a: _is_action_legal(board, a))(actions)
return jnp.zeros(26 * 6, dtype=jnp.bool_).at[actions].set(legal)

Behaviour & perf

Pure refactor, no behavioural change. Verified outputs match the previous implementation across 400 boards × 6 dice and the full _legal_action_mask over 4 dice-sets (0 mismatches). The full _legal_action_mask — called on every step (_init, _update_by_action, _change_turn) — is ~2× faster (0.61 → 0.31 ms, B=256, CPU). tests/test_backgammon.py: 20 passed (the one test_api failure is pre-existing and reproduces unchanged on main — a JAX-version issue unrelated to this change).