code: def _shuffle_jax(
      value, seed=None, name=None
  ):  # pylint: disable=unused-argument
      import jax.random as jaxrand  # pylint: disable=g-import-not-at-top
  
      if seed is None:
          raise ValueError("Must provide PRNGKey to sample in JAX.")
      return jaxrand.shuffle(seed, value, axis=0)
  
imports: import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
label_desc: Use image node_id 13 for calling a global function with example usage: _shuffle_jax(value, seed, name) and returns: jaxrand