code: def test_partition_n(self, n):
      cnt = itertools.count()
      fn = lambda m, n, v: next(cnt)
      structure = {f"layer_{i}": {"w": None} for i in range(n)}
      structures = filtering.partition_n(fn, structure, n)
      self.assertLen(structures, n)
      self.assertEqual(filtering.merge(*structures), structure)
      for i, substructure in enumerate(structures):
          expected = {f"layer_{i}": {"w": None}}
          self.assertEqual(substructure, expected)
  
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 2 for calling the FilteringTest obj's underlying member method code with example usage: obj.test_partition_n(n) without return types