code: def jax_fn_with_filter(
      jax_fn: Callable[..., Any],
      f: Callable[..., Any],
      predicate: Callable[[str, str, jax.Array], bool],
      **jax_fn_kwargs,
  ) -> Callable[..., Any]:
      """Applies a jax functionn to a given function after modifying its signature.
  
      `jax_fn_with_filter` operates in two steps:
      1   it wraps the input function `f`, which is expect to take as first
          argument a `Params` data structure, with a function taking as first two
          inputs a bipartition of the orginal parameters
      2   the resulting function is transformed with `jax_fn` and wrapped
          by a function supporting `f`'s signature and taking care of partitioning
          the `f`'s `Params` input using `predicate`.
  
      Args:
        jax_fn: jax function, e.g. `jax.grad` or `jax.jacobian`.
        f: callable to be transformed.
        predicate: predicate to be used to partition `f`'s input parameters.
        **jax_fn_kwargs: kwargs to be forwarded to `jax_fn`.
  
      Returns:
        Function calling the input jax function on the wrapped `f`.
      """
  
      def wrapper(p1, p2, *args, **kwargs):
          return f(filtering.merge(p1, p2), *args, **kwargs)
  
      jaxed_fn = jax_fn(wrapper, **jax_fn_kwargs)
  
      def fn_with_filter(p, *args, **kwargs):
          p1, p2 = filtering.partition(predicate, p)
          return jaxed_fn(p1, p2, *args, **kwargs)
  
      return fn_with_filter
  
imports: import collections
  from collections.abc import Sequence
  import itertools
  import re
  import types
  from typing import Any, Callable
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import basic
  from haiku._src import data_structures
  from haiku._src import filtering
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: jax_fn_with_filter(jax_fn, f, predicate) and returns: fn_with_filter, f, jaxed_fn