code: def partition(
      pytree: PyTree,
      filter_spec: PyTree[AxisSpec],
      replace: Any = None,
      is_leaf: Optional[Callable[[Any], bool]] = None,
  ) -> tuple[PyTree, PyTree]:
      """Splits a PyTree into two pieces. Equivalent to
      `filter(...), filter(..., inverse=True)`, but slightly more efficient.
  
      !!! info
  
          See also [`equinox.combine`][] to reconstitute the PyTree again.
      """
  
      filter_tree = jtu.tree_map(
          _make_filter_tree(is_leaf), filter_spec, pytree
      )
      left = jtu.tree_map(
          lambda mask, x: x if mask else replace, filter_tree, pytree
      )
      right = jtu.tree_map(
          lambda mask, x: replace if mask else x, filter_tree, pytree
      )
      return left, right
  
file_path: equinox/equinox/_filters.py
imports: from collections.abc import Callable
  from typing import Any, Optional, Union
  import jax
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import PyTree
label_desc: Use image node_id 7 for calling a global function with example usage: partition(pytree, filter_spec, replace, is_leaf) and returns: left,  right