code: def test_pickle_pjit_lower(self):
      def fun(x):
          return x * x
  
      with jax.sharding.Mesh(np.array(jax.devices()), ("data",)):
          lowered = pjit(
              fun, in_shardings=P("data"), out_shardings=P(None, "data")
          ).lower(core.ShapedArray(shape=(8, 8), dtype=np.float32))
  
      def verify_serialization(lowered):
          serialized, in_tree, out_tree = serialize(lowered.compile())
          compiled = deserialize_and_load(serialized, in_tree, out_tree)
          self.assertEqual(
              compiled.as_text(), lowered.compile().as_text()
          )
  
      verify_serialization(lowered)
      verify_serialization(
          jax.jit(lambda x: x * x).lower(np.arange(100))
      )
      verify_serialization(
          jax.pmap(lambda x: x * x).lower(
              np.zeros((len(jax.devices()), 4), dtype=np.float32)
          )
      )
  
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 1 for calling the JaxAotTest obj's underlying member method code with example usage: obj.test_pickle_pjit_lower() and returns: unknown