class_docstr: Compatibility tests that use TF.
  
  Uses tf.Graph to serialize and run the functions; expects that `func`
  contains a `jax2tf.call_tf` and uses `jax2tf.convert` to generate a
  `tf.Graph` containing a XlaCallModule with the actual MLIR module.
class_docstr_tok: ["Compatibility","tests","that","use","TF",".","Uses","tf.Graph","to","serialize","and","run","the","functions",";","expects","that","`","func","`","contains","a","`","jax2tf.call_tf","`","and","uses","`","jax2tf.convert","`","to","generate","a","`","tf.Graph","`","containing","a","XlaCallModule","with","the","actual","MLIR","module","."]
code: def serialize(
      self,
      func: Callable,
      data: bctu.CompatTestData,
      polymorphic_shapes: Sequence[str] | None = None,
      allow_unstable_custom_call_targets: Sequence[str] = (),
  ):
      # We serialize as a tf.Graph
      assert len(data.inputs) == 1  # We only support a single input now
      tf_graph = self.tf_func.get_concrete_function(*data.inputs).graph
      for op in tf_graph.get_operations():
          if op.type == "XlaCallModule":
              serialized_module = op.get_attr("module")
              module_str = (
                  xla_extension.mlir.deserialize_portable_artifact(
                      serialized_module
                  )
              )
              module_version = op.get_attr("version")
              break
      else:
          raise ValueError("Cannot find an XlaCallModule")
      tf_graph_def = tf_graph.as_graph_def()
      # module_str is just for human readability, add both the MLIR module
      # and the tf.Graph
      module_str = (
          "# First the MLIR module:\n"
          + module_str
          + "\n# Then the tf.Graph:\n"
          + str(tf_graph_def)
      )
      # serialized = tf_graph_def.SerializeToString()
      module = tf.Module()
      module.call = self.tf_func.get_concrete_function(*data.inputs)
      root_dir = self.create_tempdir()
      saved_model_dir = os.path.join(root_dir, "saved_model")
      os.mkdir(saved_model_dir)
      tf.saved_model.save(
          module,
          saved_model_dir,
          options=tf.saved_model.SaveOptions(
              experimental_custom_gradients=False
          ),
      )
      serialized = serialize_directory(saved_model_dir)
      nr_devices = 1
      return serialized, module_str, module_version, nr_devices
  
imports: from __future__ import annotations
  import base64
  from collections.abc import Sequence
  import io
  import os
  import tarfile
  from typing import Callable, Optional
  from absl.testing import absltest
  from jax import config
  from jax._src import test_util
  from jax._src.internal_test_util import export_back_compat_test_util
  from jax._src.lib import xla_extension
  from jax.experimental import jax2tf
  from jax.experimental.jax2tf.tests.back_compat_testdata import tf_call_tf_function
  import jax.numpy
  import tensorflow
label_desc: Use image node_id 2 for calling the CompatTensoflowTest obj's underlying member method code with example usage: obj.serialize(func, data, polymorphic_shapes, allow_unstable_custom_call_targets) and returns: serialized,  module_str,  module_version,  nr_devices