code: def cached_filter_eval_shape(fn, *args, **kwargs):
      tree = jtu.tree_map(_to_struct, (fn, args, kwargs))
      leaves, treedef = jtu.tree_flatten(tree)
      leaves = tuple(leaves)
      return _cached_filter_eval_shape(leaves, treedef)
  
file_path: equinox/equinox/_eval_shape.py
imports: import functools
  from collections.abc import Callable
  from typing import Any, Union
  import jax
  import jax._src.traceback_util
  import jax.tree_util
  from jaxtyping import PyTree
  from ._caches import internal_lru_caches
  from ._filters import combine, is_array, partition
  from ._module import Static
label_desc: Use image node_id 5 for calling a global function with example usage: cached_filter_eval_shape(fn) and returns: _cached_filter_eval_shape