code: def test_topology_pjit_serialize(self):
      try:
          aot_topo = topologies.get_topology_desc(
              platform=jax.devices()[0].platform
          )
      except NotImplementedError:
          raise unittest.SkipTest("PJRT Topology not supported")
  
      if jtu.TEST_WITH_PERSISTENT_COMPILATION_CACHE.value:
          raise unittest.SkipTest(
              "Compilation caching not yet supported."
          )
  
      @jax.jit
      def fn(x):
          return x * x
  
      def lower_and_load(mesh):
          s = jax.sharding.NamedSharding(mesh, P("x", "y"))
          x_shape = jax.ShapeDtypeStruct(
              shape=(16, 16), dtype=jnp.dtype("float32"), sharding=s
          )
          lowered = fn.lower(x_shape)
          serialized, in_tree, out_tree = serialize(lowered.compile())
          compiled = deserialize_and_load(serialized, in_tree, out_tree)
          return compiled
  
      ref_topo = topologies.get_attached_topology()
      n = max(1, len(ref_topo.devices) // 2)
      mesh_shape = (len(ref_topo.devices) // n, n)
  
      ref_mesh = topologies.make_mesh(ref_topo, mesh_shape, ("x", "y"))
      aot_mesh = topologies.make_mesh(aot_topo, mesh_shape, ("x", "y"))
      self.assertEqual(
          lower_and_load(ref_mesh).as_text(),
          lower_and_load(aot_mesh).as_text(),
      )
  
file_path: jax/tests/aot_test.py
imports: import contextlib
  import unittest
  from absl.testing import absltest
  import jax
  from jax import config
  from jax._src import core
  from jax._src import test_util
  from jax._src.lib import xla_client
  from jax.experimental import topologies
  from jax.experimental.pjit import pjit
  from jax.experimental.serialize_executable import deserialize_and_load, serialize
  import jax.numpy
  from jax.sharding import PartitionSpec
  import numpy
label_desc: Use image node_id 2 for calling the JaxAotTest obj's underlying member method code with example usage: obj.test_topology_pjit_serialize() and returns: unknown, compiled