def _mask_tree(tree: search.Tree, num_actions: int, mode: str) -> search.Tree:
  """Masks out parts of the tree based upon node type.

  "Actions" in our tree can either be action or chance values: A' = A + C. This
  utility function masks the parts of the tree containing dimensions of shape
  A' to be either A or C depending upon `mode`.

  Args:
    tree: The tree to be masked.
    num_actions: The number of environment actions A.
    mode: Either "decision" or "chance".

  Returns:
    An appropriately masked tree.
  """

  def _take_slice(x):
    if mode == 'decision':
      return x[..., :num_actions]
    elif mode == 'chance':
      return x[..., num_actions:]
    else:
      raise ValueError(f'Unknown mode: {mode}.')

  return tree.replace(
      children_index=_take_slice(tree.children_index),
      children_prior_logits=_take_slice(tree.children_prior_logits),
      children_visits=_take_slice(tree.children_visits),
      children_rewards=_take_slice(tree.children_rewards),
      children_discounts=_take_slice(tree.children_discounts),
      children_values=_take_slice(tree.children_values),
      root_invalid_actions=_take_slice(tree.root_invalid_actions))
Masks out parts of the tree based upon node type.

  "Actions" in our tree can either be action or chance values: A' = A + C. This
  utility function masks the parts of the tree containing dimensions of shape
  A' to be either A or C depending upon `mode`.

  Args:
    tree: The tree to be masked.
    num_actions: The number of environment actions A.
    mode: Either "decision" or "chance".

  Returns:
    An appropriately masked tree.
  
mctx/_src/policies.py