docstr: Compares jax_func(*args) with convert(jax_func)(*args).
  
  It compares the result of JAX, TF ("eager" mode),
  TF with tf.function ("graph" mode), and TF with
  tf.function(jit_compile=True) ("compiled" mode). In each mode,
  either we expect to encounter a known limitation, or the value should
  match the value from the JAX execution.
  
  Args:
    func_jax: the function to invoke (``func_jax(*args)``)
    args: the arguments.
    enable_xla: if True, allows the use of XLA ops in jax2tf.convert
      (default: True).
    limitations: the set of limitations for this harness (not yet filtered
      by mode).
code: def ConvertAndCompare(
      self,
      func_jax: Callable,
      *args,
      enable_xla: bool = True,
      limitations: Sequence = (),
  ):
      """Compares jax_func(*args) with convert(jax_func)(*args).
  
      It compares the result of JAX, TF ("eager" mode),
      TF with tf.function ("graph" mode), and TF with
      tf.function(jit_compile=True) ("compiled" mode). In each mode,
      either we expect to encounter a known limitation, or the value should
      match the value from the JAX execution.
  
      Args:
        func_jax: the function to invoke (``func_jax(*args)``)
        args: the arguments.
        enable_xla: if True, allows the use of XLA ops in jax2tf.convert
          (default: True).
        limitations: the set of limitations for this harness (not yet filtered
          by mode).
      """
      # Run JAX. Should not fail, we assume that the harness has been filtered
      # already by JAX unimplemented primitives.
      result_jax = func_jax(*args)  # JAX
      result_tf = None
  
      func_tf = jax2tf.convert(func_jax, enable_xla=enable_xla)
  
      unexpected_successes: list[str] = []
      # Run the "compiled" mode first, it is most important
      for mode in ("compiled", "eager", "graph"):
          if mode == "graph" and jtu.device_under_test() == "tpu":
              continue  # The "graph" mode on TPU is the same as "compiled"
  
          def log_message(extra):
              return f"[{self._testMethodName}] {mode=}: {extra}"
  
          jax2tf_limits = tuple(
              filter(lambda l: l.filter(mode=mode), limitations)
          )
  
          skip_tf_run = [l for l in jax2tf_limits if l.skip_tf_run]
          if skip_tf_run:
              logging.info(
                  log_message(
                      f"Skip TF run due to limitations {skip_tf_run}"
                  )
              )
              continue
  
          try:
              result_tf = _run_tf_function(func_tf, *args, mode=mode)
              tf_exception = None
          except Exception as e:
              tf_exception = e
  
          expect_tf_error = [
              l for l in jax2tf_limits if l.expect_tf_error
          ]
          if tf_exception:
              if expect_tf_error:
                  logging.info(
                      log_message(
                          "Found expected TF error with enabled limitations "
                          f"{expect_tf_error}; TF error is {tf_exception}"
                      )
                  )
                  continue
              else:
                  raise tf_exception
          else:
              if expect_tf_error:
                  # It is more ergonomic to print all successful modes once
                  logging.warning(
                      log_message(
                          f"Unexpected execution success with known limitations {expect_tf_error}"
                      )
                  )
                  unexpected_successes.append(
                      f"{mode}: {expect_tf_error}"
                  )
  
          if (
              jtu.device_under_test() == "gpu"
              and "dot_general_preferred" in self._testMethodName
          ):
              logging.info(
                  log_message(
                      f"Arguments are {args}, JAX result is {result_jax}\nand TF result is {result_tf}"
                  )
              )
  
          skip_comparison = [
              l for l in jax2tf_limits if l.skip_comparison
          ]
          if skip_comparison:
              logging.warning(
                  log_message(
                      f"Skip result comparison due to {skip_comparison}"
                  )
              )
              continue
  
          max_tol = None
          max_tol_lim = (
              None
              if not jax2tf_limits
              else jax2tf_limits[0].get_max_tolerance_limitation(
                  jax2tf_limits
              )
          )
          if max_tol_lim is not None:
              max_tol = max_tol_lim.tol
              logging.info(
                  log_message(
                      f"Using tol={max_tol} due to {max_tol_lim}"
                  )
              )
  
          # Convert results to np.arrays
          result_tf = tf.nest.map_structure(lambda t: t.numpy(), result_tf)  # type: ignore
  
          custom_assert_lim = [
              l for l in jax2tf_limits if l.custom_assert
          ]
          assert (
              len(custom_assert_lim) <= 1
          ), f"Expecting at most one applicable limitation with custom_assert, found {custom_assert_lim}"
  
          try:
              err_msg = f"TF mode {mode}."
              log_hlo_on_error = (
                  mode == "compiled" or jtu.device_under_test() == "tpu"
              )
              if log_hlo_on_error:
                  err_msg += (
                      " See the logs for JAX and TF HLO comparisons."
                  )
              if custom_assert_lim:
                  logging.info(
                      log_message(
                          f"Running custom_assert with tol={max_tol} due to {custom_assert_lim[0]}"
                      )
                  )
                  custom_assert_lim[0].custom_assert(
                      self,
                      result_jax,
                      result_tf,
                      args=args,
                      tol=max_tol,
                      err_msg=err_msg,
                  )
              else:
                  logging.info(
                      log_message(
                          f"Running default assert with tol={max_tol}"
                      )
                  )
                  self.assertAllClose(
                      result_jax,
                      result_tf,
                      atol=max_tol,
                      rtol=max_tol,
                      err_msg=err_msg,
                  )
          except AssertionError as e:
              # Print the HLO for comparison
              if not log_hlo_on_error:
                  print(
                      f"[{self._testMethodName}] Not logging HLO because the "
                      f"mode was {mode}"
                  )
                  raise
  
              logging.info(
                  "[%s] Logging HLO for exception in mode %s: %s",
                  self._testMethodName,
                  mode,
                  e,
              )
              jax_lowered = jax.jit(func_jax).lower(*args)
              # We log the HLO dialect for easier comparison with TF
              logging.info(
                  "[%s] JAX NON_OPT HLO\n%s",
                  self._testMethodName,
                  jax_lowered.compiler_ir(dialect="hlo").as_hlo_text(),
              )  # type: ignore
  
              tf_args_signature = _make_tf_input_signature(*args)
              # If we give the signature, we cannot pass scalars
              tf_args_no_scalars = tuple(
                  map(
                      lambda a, sig: tf.convert_to_tensor(
                          a, dtype=sig.dtype
                      ),
                      args,
                      tf_args_signature,
                  )
              )
  
              tf_func_compiled = tf.function(
                  func_tf,
                  autograph=False,
                  jit_compile=True,
                  input_signature=tf_args_signature,
              )
              tf_hlo = tf_func_compiled.experimental_get_compiler_ir(
                  *tf_args_no_scalars
              )(stage="hlo")
              logging.info(
                  "[%s] TF NON OPT HLO\n{%s}",
                  self._testMethodName,
                  tf_hlo,
              )
  
              backend = xla_bridge.get_backend()
              modules = backend.compile(
                  str(jax_lowered.compiler_ir())
              ).hlo_modules()
              jax_opt_hlo = modules[0].to_string()
              logging.info(
                  "[%s] JAX OPT HLO\n%s",
                  self._testMethodName,
                  jax_opt_hlo,
              )
  
              tf_opt_hlo = (
                  tf_func_compiled.experimental_get_compiler_ir(
                      *tf_args_no_scalars
                  )(stage="optimized_hlo")
              )
              logging.info(
                  "[%s] TF OPT HLO\n%s",
                  self._testMethodName,
                  tf_opt_hlo,
              )
  
              raise
  
      # end "for mode"
  
      if unexpected_successes:
          msg = (
              f"[{self._testMethodName}] The following are unexpected "
              "successful modes:\n" + "\n".join(unexpected_successes)
          )
          logging.warning(msg)
          # Uncomment the below if you want to see warnings as failures
          # self.assertEmpty(msg)
      return result_jax, result_tf
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import re
  import os
  from typing import Any, Callable, Optional
  from absl.testing import absltest
  from absl import logging
  import jax
  from jax import dtypes
  from jax import numpy
  from jax._src import test_util
  from jax import tree_util
  from jax.experimental import jax2tf
  from jax.experimental import export
  from jax._src import config
  from jax._src import xla_bridge
  import numpy
  import tensorflow
  from tensorflow.compiler.xla import xla_data_pb2
  from tensorflow.compiler.tf2xla.python import xla
label_desc: Use image node_id 3 for calling the JaxToTfTestCase obj's underlying member method code with example usage: obj.ConvertAndCompare(func_jax) and returns: result_jax,  result_tf, str, self, str, mode, str, extra