code: def tree_select(pred: Array, on_true: Any, on_false: Any):
      """Select either one of two identical nested structs based on condition.
  
      Args:
        pred: a boolean condition.
        on_true: an arbitrary nested structure.
        on_false: a nested structure identical to `on_true`.
  
      Returns:
        the selected nested structure.
      """
      if tree_structure(on_true) != tree_structure(on_false):
          raise ValueError(
              "The two branches must have the same structure."
          )
      return jax.tree_util.tree_map(
          lambda x, y: jax.lax.select(pred, x, y), on_true, on_false
      )
  
file_path: rlax/rlax/_src/tree_util.py
imports: import functools
  from typing import Any, Callable, Sequence
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 1 for calling a global function with example usage: tree_select(pred, on_true, on_false) and returns: jax