code: def test_transforms_with_filter(self):
      # Note to make sense of test:
      #
      # out = (w0 + b0) * w1 + b1
      #     = w0 * w1 + b0 * w1 + b1
      # doutdw0 = w1
      # doutdw1 = w0 + b0
      # with w0 = 1.0, b0 = 1.5, w1 = 3.0, b1 = 4.5
      init_fn, apply_fn = transform.transform(get_net)
      inputs = jnp.ones((1, 1))
      params = init_fn(jax.random.PRNGKey(428), inputs)
  
      df_fn = jax_fn_with_filter(
          jax_fn=jax.grad,
          f=apply_fn,
          predicate=lambda module_name, name, _: name == "w",
      )
      df = df_fn(params, None, inputs)
      self.assertEqual(
          to_set(df),
          {("first_layer/w", (3.0,)), ("second_layer/w", (2.5,))},
      )
  
      fn = jax_fn_with_filter(
          jax_fn=jax.value_and_grad,
          f=apply_fn,
          predicate=lambda module_name, name, _: name == "w",
      )
      v = fn(params, None, inputs)
      self.assertEqual(v[0], jnp.array([12.0]))
      self.assertEqual(to_set(df), to_set(v[1]))
  
      def get_stacked_net(x):
          y = get_net(x)
          return jnp.stack([y, 2.0 * y])
  
      _, apply_fn = transform.transform(get_stacked_net)
      jf_fn = jax_fn_with_filter(
          jax_fn=jax.jacobian,
          f=apply_fn,
          predicate=lambda module_name, name, _: name == "w",
      )
      jf = jf_fn(params, None, inputs)
  
      self.assertEqual(
          to_set(jf),
          {
              ("first_layer/w", (3.0, 6.0)),
              ("second_layer/w", (2.5, 5.0)),
          },
      )
  
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 10 for calling the FilteringTest obj's underlying member method code with example usage: obj.test_transforms_with_filter() and returns: jnp