code: def testPickleOfKeyArray(self, prng_name):
      with jax.default_prng_impl(prng_name):
          k1 = jax.random.PRNGKey(72)
          s = pickle.dumps(k1)
          k2 = pickle.loads(s)
          self.assertEqual(k1.dtype, k2.dtype)
          with jax.legacy_prng_key("allow"):
              self.assertArraysEqual(
                  jax.random.key_data(k1), jax.random.key_data(k2)
              )
  
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 3 for calling the PickleTest obj's underlying member method code with example usage: obj.testPickleOfKeyArray(prng_name) without return types