code: def assert_output_type(self, out_cls):
      def assert_type_recursive(s):
          self.assertEqual(type(s), out_cls)
  
      for in_cls in (dict, data_structures.FlatMap):
          with self.subTest(str(in_cls)):
              structure_a = in_cls({"m1": in_cls({"w": None})})
              structure_b = in_cls({"m2": in_cls({"w": None})})
              structure_c = in_cls(
                  {f"{i}": in_cls({"w": None}) for i in range(5)}
              )
              assert_type_recursive(
                  filtering.filter(lambda m, n, v: True, structure_a)
              )
              assert_type_recursive(
                  filtering.map(lambda m, n, v: v, structure_a)
              )
              assert_type_recursive(
                  filtering.merge(structure_a, structure_b)
              )
              parts = filtering.partition(
                  lambda m, n, v: int(m) > 1, structure_c
              )
              for part in parts:
                  assert_type_recursive(part)
              parts = filtering.partition_n(
                  lambda m, n, v: int(m), structure_c, 5
              )
              for part in parts:
                  assert_type_recursive(part)
  
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 16 for calling the FilteringTest obj's underlying member method code with example usage: obj.assert_output_type(out_cls) without return types