code: def test_tree_functions(self):
      f = FlatMap({"foo": {"b": {"c": 1}, "d": 2}, "bar": {"c": 1}})
  
      m = jax.tree_util.tree_map(lambda x: x + 1, f)
      self.assertEqual(type(m), FlatMap)
      self.assertEqual(
          m, {"foo": {"b": {"c": 2}, "d": 3}, "bar": {"c": 2}}
      )
  
      mm = jax.tree_util.tree_map(lambda x, y: x + y, f, f)
      self.assertEqual(type(mm), FlatMap)
      self.assertEqual(
          mm, {"foo": {"b": {"c": 2}, "d": 4}, "bar": {"c": 2}}
      )
  
      leaves, treedef = jax.tree_util.tree_flatten(f)
      self.assertEqual(leaves, [1, 1, 2])
      uf = jax.tree_util.tree_unflatten(treedef, leaves)
      self.assertEqual(type(f), FlatMap)
      self.assertEqual(f, uf)
  
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 24 for calling the FlatMappingTest obj's underlying member method code with example usage: obj.test_tree_functions() without return types