def _produce_prediction_output(rng_key, num_actions):
  """Producing the model output as in the stored expected trees."""
  policy_rng, value_rng, reward_rng = jax.random.split(rng_key, 3)
  policy_rng, value_rng, reward_rng = jax.random.split(rng_key, 3)
  del rng_key
  # Producing value from [-1, +1).
  value = jax.random.uniform(value_rng, shape=(), minval=-1.0, maxval=1.0)
  # Producing reward from [-1, +1).
  reward = jax.random.uniform(reward_rng, shape=(), minval=-1.0, maxval=1.0)
  return dict(
      policy_logits=jax.random.normal(policy_rng, shape=[num_actions]),
      value=value,
      reward=reward,
  )
Producing the model output as in the stored expected trees.
mctx/_src/tests/tree_test.py