def tree_to_pytree(tree: mctx.Tree, batch_i: int = 0):
  """Converts the MCTS tree to nested dicts."""
  nodes = {}
  nodes[0] = _create_pynode(
      tree, batch_i, 0, prior=1.0, action=None, reward=None)
  children_prior_probs = jax.nn.softmax(tree.children_prior_logits, axis=-1)
  for node_i in range(tree.num_simulations + 1):
    for a_i in range(tree.num_actions):
      prior = children_prior_probs[batch_i, node_i, a_i]
      # Index of children, or -1 if not expanded
      child_i = int(tree.children_index[batch_i, node_i, a_i])
      if child_i >= 0:
        reward = tree.children_rewards[batch_i, node_i, a_i]
        child = _create_pynode(
            tree, batch_i, child_i, prior=prior, action=a_i, reward=reward)
        nodes[child_i] = child
      else:
        child = _create_bare_pynode(prior=prior, action=a_i)
      # pylint: disable=line-too-long
      nodes[node_i]["child_stats"].append(child)  # pytype: disable=attribute-error
      # pylint: enable=line-too-long
  return nodes[0]
Converts the MCTS tree to nested dicts.
mctx/_src/tests/tree_test.py