code: def setUp(self):
      super().setUp()
      # Ensure that all TF ops are created on the proper device (TPU or GPU or CPU)
      tf_preferred_devices = (
          tf.config.list_logical_devices("TPU")
          + tf.config.list_logical_devices("GPU")
          + tf.config.list_logical_devices()
      )
      self.tf_default_device = tf_preferred_devices[0]
      logging.info(
          "Running jax2tf converted code on %s.", self.tf_default_device
      )
      # We need --config=cuda build flag for TF to see the GPUs
      self.assertEqual(
          jtu.device_under_test().upper(),
          self.tf_default_device.device_type,
      )
  
      # We run the tests using the maximum version supported, even though
      # the default serialization version may be held back for a while to
      # ensure compatibility
      version = config.jax_serialization_version.value
      if self.use_max_serialization_version:
          # Use the largest supported by both export and tfxla.call_module
          version = min(
              export.maximum_supported_serialization_version,
              tfxla.call_module_maximum_supported_version(),
          )
          self.assertGreaterEqual(
              version, export.minimum_supported_serialization_version
          )
          self.enter_context(config.jax_serialization_version(version))
      logging.info(
          "Using JAX serialization version %s (export.max_version %s, tf.XlaCallModule max version %s)",
          version,
          export.maximum_supported_serialization_version,
          tfxla.call_module_maximum_supported_version(),
      )
  
      with contextlib.ExitStack() as stack:
          stack.enter_context(tf.device(self.tf_default_device))
          self.addCleanup(stack.pop_all().close)
  
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 1 for calling the JaxToTfTestCase obj's underlying member method code with example usage: obj.setUp() without return types