code: def filter_primitive_def(rule):
      """For wrapping def_impl and def_abstract_eval.
  
      These can now take arbitrary inputs and outputs.
      """
  
      def _wrapper(*dynamic, treedef, static, flatten):
          args = jtu.tree_unflatten(treedef, _combine(dynamic, static))
          out = rule(*args)
          return flatten(out)
  
      return _wrapper
  
imports: import functools
  from typing import Any
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.tree_util
  from jaxtyping import PyTree
  from .._filters import combine, is_array, is_array_like, partition
  from .._tree import tree_equal
label_desc: Use image node_id 10 for calling a global function with example usage: filter_primitive_def(rule) and returns: _wrapper, flatten