code: def filter(
      pytree: PyTree,
      filter_spec: PyTree[AxisSpec],
      inverse: bool = False,
      replace: Any = None,
      is_leaf: Optional[Callable[[Any], bool]] = None,
  ) -> PyTree:
      """
      Filters out the leaves of a PyTree not satisfying a condition. Those not satisfying
      the condition are replaced with `replace`.
  
      !!! Example
  
          ```python
          pytree = [(jnp.array(0), 1), object()]
          result = eqx.filter(pytree, eqx.is_array)
          # [(jnp.array(0), None), None]
          ```
  
      !!! Example
  
          ```python
          pytree = [(jnp.array(0), 1), object()]
          result = eqx.filter(pytree, [(False, False), True])
          # [(None, None), object()]
          ```
  
      **Arguments:**
  
      - `pytree` is any PyTree.
      - `filter_spec` is a PyTree whose structure should be a prefix of the structure of
          `pytree`. Each of its leaves should either be:
          - `True`, in which case the leaf or subtree is kept;
          - `False`, in which case the leaf or subtree is replaced with `replace`;
          - a callable `Leaf -> bool`, in which case this is evaluted on the leaf or
              mapped over the subtree, and the leaf kept or replaced as appropriate.
      - `inverse` switches the truthy/falsey behaviour: falsey results are kept and
          truthy results are replaced.
      - `replace` is what to replace any falsey leaves with. Defaults to `None`.
      - `is_leaf`: Optional function called at each node of the PyTree. It should return
          a boolean. `True` indicates that the whole subtree should be treated as leaf;
          `False` indicates that the subtree should be traversed as a PyTree.
  
      **Returns:**
  
      A PyTree of the same structure as `pytree`.
      """
  
      inverse = bool(
          inverse
      )  # just in case, to make the != trick below work reliably
      filter_tree = jtu.tree_map(
          _make_filter_tree(is_leaf), filter_spec, pytree
      )
      return jtu.tree_map(
          lambda mask, x: x if bool(mask) != inverse else replace,
          filter_tree,
          pytree,
      )
  
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 6 for calling a global function with example usage: filter(pytree, filter_spec, inverse, replace, is_leaf) and returns: jtu