code: def testPickleOfPmappedFunctions(self):
      @jax.pmap
      def f(x, y):
          return x * y
  
      @jax.pmap
      def g(z):
          return f(z, z + 77)  # noqa: F821
  
      expected = g(jnp.asarray([[32]]))
      s = cloudpickle.dumps(g)
      del f, g
  
      g_unpickled = pickle.loads(s)
      actual = g_unpickled(jnp.asarray([[32]]))
      self.assertEqual(expected, actual)
  
imports: import pickle
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import numpy
  from jax import config
  from jax.interpreters import pxla
  from jax._src import test_util
  from jax._src.lib import xla_client
  import numpy
label_desc: Use image node_id 2 for calling the CloudpickleTest obj's underlying member method code with example usage: obj.testPickleOfPmappedFunctions() and returns: unknown, f