def _fold_action_in(rng_key, action, num_actions):
  """Returns a new rng key, selected by the given action."""
  chex.assert_shape(action, ())
  chex.assert_type(action, jnp.int32)
  sub_rngs = jax.random.split(rng_key, num_actions)
  return sub_rngs[action]
Returns a new rng key, selected by the given action.
mctx/_src/tests/tree_test.py