code: def test_frozen_builtins_jax_compatibility(self):
      f = FlatMap({"foo": [3, 2], "bar": {"a": 3}})
      mapped_frozen_list = jax.tree_util.tree_map(
          lambda x: x + 1, f["foo"]
      )
      self.assertEqual(mapped_frozen_list[0], 4)
  
      mapped_frozen_dict = jax.tree_util.tree_map(
          lambda x: x + 1, f["bar"]
      )
      self.assertEqual(mapped_frozen_dict["a"], 4)
  
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 29 for calling the FlatMappingTest obj's underlying member method code with example usage: obj.test_frozen_builtins_jax_compatibility() without return types