def _prepare_recurrent_fn(num_actions, *, discount, zero_reward):
  """Returns a dynamics function consistent with the expected trees."""

  def recurrent_fn(params, rng_key, action, embedding):
    del params, rng_key
    # The embeddings serve as rng_keys.
    embedding = jax.vmap(
        functools.partial(_fold_action_in, num_actions=num_actions))(embedding,
                                                                     action)
    output = jax.vmap(
        functools.partial(_produce_prediction_output, num_actions=num_actions))(
            embedding)
    reward = output["reward"]
    if zero_reward:
      reward = jnp.zeros_like(reward)
    return mctx.RecurrentFnOutput(
        reward=reward,
        discount=jnp.full_like(reward, discount),
        prior_logits=output["policy_logits"],
        value=output["value"],
    ), embedding

  return recurrent_fn
Returns a dynamics function consistent with the expected trees.
mctx/_src/tests/tree_test.py