code: def write_and_check_harness(
      harness: primitive_harness.Harness,
      io: Io,
      save_directory: str,
      for_platforms: Sequence[str] = ("cpu", "tpu"),
  ) -> Sequence[str]:
      """Writes and checks HLO for a given harness.
  
      Writes the HLOs generated in the current platform for all platforms.
      If it finds previously written HLOs generated on other platforms, compares
      them with the ones generated on this platform.
  
      Returns a list of harnesses on which diffs were found.
      """
      diffs = []
  
      func_jax = harness.dyn_fun
      rng = npr.RandomState(zlib.adler32(harness.fullname.encode()))
      args = harness.dyn_args_maker(rng)
  
      # Generate the HLO for all platforms
      for for_platform in for_platforms:
          if not harness.filter(for_platform):
              logging.info(
                  "Skip harness %s for %s because it is not implemented in JAX",
                  harness.fullname,
                  for_platform,
              )
              continue
  
          scenario1 = Scenario(
              harness, jax.default_backend(), for_platform
          )
          output_file = scenario1.output_file(save_directory)
          output_dir = os.path.dirname(output_file)
          if not io.exists(output_dir):
              io.makedirs(output_dir)
  
          if io.exists(output_file):
              with open(output_file) as f:
                  hlo = f.read()
          else:
              # For a tighter check, detect the native platform lowering and do not
              # trigger cross-lowering
              if for_platform == jax.default_backend():
                  lowered = jax.jit(func_jax).lower(*args)
              else:
                  # TODO: replace this with JAX cross-platform API, without going through
                  # jax2tf
                  from jax.experimental.jax2tf.jax2tf import (
                      cross_platform_lowering,
                  )
  
                  lowered = cross_platform_lowering(
                      func_jax, args, platforms=[for_platform]
                  )
              hlo = lowered.compiler_ir(dialect="stablehlo")  # type: ignore
              with open(output_file, "w") as f:
                  f.write(str(hlo))
  
          # Compare with previously written files
          for on_platform in ["cpu", "tpu"]:
              if on_platform == jax.default_backend():
                  continue
              scenario2 = Scenario(harness, on_platform, for_platform)
              other_file = scenario2.output_file(save_directory)
              if io.exists(other_file):
                  logging.info(
                      "Comparing for %s harness %s on %s vs %s",
                      for_platform,
                      harness.fullname,
                      jax.default_backend(),
                      on_platform,
                  )
                  with open(other_file) as f:
                      other_hlo = f.read()
  
                  if hlo != other_hlo:
                      logging.info(
                          "Found diff",
                          for_platform,
                          harness.fullname,
                          jax.default_backend(),
                          on_platform,
                      )
                      diffs.append(
                          f"Found diff between {output_file} and {other_file}"
                      )
  
      return diffs
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import os
  import re
  from typing import Callable, Optional
  import zlib
  from absl import app
  from absl import logging
  import numpy.random
  import jax
  from jax import config
  from jax.experimental import jax2tf
  from jax._src import test_util
  from jax.experimental.jax2tf.tests import primitive_harness
label_desc: Use image node_id 1 for calling a global function with example usage: write_and_check_harness(harness, io, save_directory, for_platforms) and returns: diffs