code: def test_flatten_nested_struct(self):
      d = {
          "foo": {"bar": [1, 2, 3]},
          "baz": {"bat": [4, 5, 6], "qux": [7, [8, 9]]},
      }
      f = FlatMap(d)
      leaves, treedef = jax.tree_util.tree_flatten(f)
      self.assertEqual([4, 5, 6, 7, 8, 9, 1, 2, 3], leaves)
      g = jax.tree_util.tree_unflatten(treedef, leaves)
      self.assertEqual(g, f)
      self.assertEqual(g, d)
  
imports: import collections
  import copy
  import pickle
  import threading
  from absl.testing import absltest
  from absl.testing import parameterized
  import cloudpickle
  import dill
  from haiku._src import data_structures
  import jax
  import tree
label_desc: Use image node_id 25 for calling the FlatMappingTest obj's underlying member method code with example usage: obj.test_flatten_nested_struct() without return types