def _create_pynode(tree, batch_i, node_i, prior, action, reward):
  """Returns a dict with extracted search statistics."""
  node = dict(
      prior=_round_float(prior),
      visit=int(tree.node_visits[batch_i, node_i]),
      value_view=_round_float(tree.node_values[batch_i, node_i]),
      raw_value_view=_round_float(tree.raw_values[batch_i, node_i]),
      child_stats=[],
      evaluation_index=node_i,
  )
  if action is not None:
    node["action"] = action
  if reward is not None:
    node["reward"] = _round_float(reward)
  return node
Returns a dict with extracted search statistics.
mctx/_src/tests/tree_test.py