code: def tree_split_leaves(
      tree_like: Any, axis: int = 0, keepdim: bool = False
  ):
      """Splits a tree of arrays into an array of trees avoiding data copying.
  
      Note: `jax.numpy.DeviceArray`'s data gets copied.
  
      Args:
        tree_like: a nested object with leaves to split.
        axis: an axis for splitting.
        keepdim: a bool indicating whether to keep `axis` dimension.
  
      Returns:
       A tuple of `size(axis)` trees containing results of splitting.
      """
  
      # Disable pylint to correctly process `np.ndarray`s.
      if len(tree_like) == 0:  # pylint: disable=g-explicit-length-test
          return tree_like
      leaves, treedef = jax.tree_util.tree_flatten(tree_like)
      axis_size = leaves[0].shape[axis]
      split_leaves = [np.split(l, axis_size, axis=axis) for l in leaves]
      ind_ = lambda x, i: x[i] if keepdim else np.squeeze(x[i], axis)
      split_trees = (
          (ind_(l, i) for l in split_leaves) for i in range(axis_size)
      )
      return tuple(
          jax.tree_util.tree_unflatten(treedef, t) for t in split_trees
      )
  
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 4 for calling a global function with example usage: tree_split_leaves(tree_like, axis, keepdim) and returns: tuple, tree_like