code: def filter_primitive_batching(rule):
      """
      The input batch axes (to the wrapped rule) will be a PyTree with the same
      structure as the input primals, with `None` for all non-JAX-arrays.
  
      The output batch axes are expected to match the output primals, with `None`
      for all non-JAX-arrays.
      """
  
      def _wrapper(dynamic, batch_axes, *, treedef, static, flatten):
          flat = _combine(dynamic, static)
          inputs = jtu.tree_unflatten(treedef, flat)
          batch_axes = [
              None if b is batching.not_mapped else b
              for b in batch_axes
          ]
          batch_axes_static = [
              x if x is _missing_dynamic else None for x in static
          ]
          batch_axes = _combine(batch_axes, batch_axes_static)
          batch_axes = jtu.tree_unflatten(treedef, batch_axes)
          out, batch_axes = rule(inputs, batch_axes)
          flat_out, flat_batch_axes = flatten(out, batch_axes)
          flat_batch_axes = [
              batching.not_mapped if b is None else b
              for b in flat_batch_axes
          ]
          return flat_out, flat_batch_axes
  
      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 13 for calling a global function with example usage: filter_primitive_batching(rule) and returns: _wrapper, flat_out,  flat_batch_axes