def _prepare_root(batch_size, num_actions):
  """Returns a root consistent with the stored expected trees."""
  rng_key = jax.random.PRNGKey(0)
  # Using a different rng_key inside each batch element.
  rng_keys = [rng_key]
  for i in range(1, batch_size):
    rng_keys.append(jax.random.fold_in(rng_key, i))
  embedding = jnp.stack(rng_keys)
  output = jax.vmap(
      functools.partial(_produce_prediction_output, num_actions=num_actions))(
          embedding)
  return mctx.RootFnOutput(
      prior_logits=output["policy_logits"],
      value=output["value"],
      embedding=embedding,
  )
Returns a root consistent with the stored expected trees.
mctx/_src/tests/tree_test.py