def _get_deepest_leaf(tree, node_index):
  """Returns `(leaf, depth)` with maximum depth and visit count.

  Args:
    tree: _unbatched_ MCTS tree state.
    node_index: the node of the inspected subtree.

  Returns:
    `(leaf, depth)` of a deepest leaf. If multiple leaves have the same depth,
    the leaf with the highest visit count is returned.
  """
  np.testing.assert_equal(len(tree.children_index.shape), 2)
  leaf = node_index
  max_found_depth = 0
  for action in range(tree.children_index.shape[-1]):
    next_node_index = tree.children_index[node_index, action]
    if next_node_index != tree.UNVISITED:
      found_leaf, found_depth = _get_deepest_leaf(tree, next_node_index)
      if ((1 + found_depth, tree.node_visits[found_leaf]) >
          (max_found_depth, tree.node_visits[leaf])):
        leaf = found_leaf
        max_found_depth = 1 + found_depth
  return leaf, max_found_depth
Returns `(leaf, depth)` with maximum depth and visit count.

  Args:
    tree: _unbatched_ MCTS tree state.
    node_index: the node of the inspected subtree.

  Returns:
    `(leaf, depth)` of a deepest leaf. If multiple leaves have the same depth,
    the leaf with the highest visit count is returned.
  
mctx/_src/tests/policies_test.py