code: def export_and_compare_to_native(
      self,
      func_jax: Callable,
      *args: jax.Array,
      unimplemented_platforms: set[str] = set(),
      skip_run_on_platforms: set[str] = set(),
      tol: float | None = None,
  ):
      devices = [
          d
          for d in self.__class__.devices
          if d.platform not in unimplemented_platforms
      ]
      logging.info("Using devices %s", [str(d) for d in devices])
      # lowering_platforms uses "cuda" instead of "gpu"
      lowering_platforms: list[str] = [
          p if p != "gpu" else "cuda"
          for p in ("cpu", "gpu", "tpu")
          if p not in unimplemented_platforms
      ]
  
      if len(lowering_platforms) <= 1:
          self.skipTest(
              "Harness is uninteresting with fewer than 2 platforms"
          )
  
      logging.info("Exporting harness for %s", lowering_platforms)
      exp = export.export(
          func_jax, lowering_platforms=lowering_platforms
      )(*args)
  
      for device in devices:
          if device.platform in skip_run_on_platforms:
              logging.info("Skipping running on %s", device)
              continue
          device_args = jax.tree_util.tree_map(
              lambda x: jax.device_put(x, device), args
          )
          logging.info("Running harness natively on %s", device)
          native_res = func_jax(*device_args)
          logging.info("Running exported harness on %s", device)
          exported_res = export.call_exported(exp)(*device_args)
          if tol is not None:
              logging.info(f"Using non-standard tolerance {tol}")
          self.assertAllClose(
              native_res, exported_res, atol=tol, rtol=tol
          )
  
file_path: jax/tests/export_harnesses_multi_platform_test.py
imports: from __future__ import annotations
  import math
  import re
  from typing import Callable
  from absl import logging
  from absl.testing import absltest
  import numpy
  import jax
  from jax import lax
  from jax._src import test_util
  from jax.experimental import export
  from jax._src.internal_test_util import test_harnesses
label_desc: Use image node_id 3 for calling the PrimitiveTest obj's underlying member method code with example usage: obj.export_and_compare_to_native(func_jax) without return types