code: def get_prng(seed):
      """JAX-compatible way of getting PRNG seeds."""
      if np.shape(seed):
          raise TypeError("PRNGKey seed must be a scalar.")
      convert = lambda k: np.reshape(np.asarray(k, np.uint32), [1])
      k1 = convert(np.bitwise_and(np.right_shift(seed, 32), 0xFFFFFFFF))
      k2 = convert(np.bitwise_and(seed, 0xFFFFFFFF))
      return np.concatenate([k1, k2], 0)
  
file_path: trax/trax/fastmath/numpy.py
imports: import numpy
  from scipy.special import logsumexp
  from trax.shapes import signature
label_desc: Use image node_id 1 for calling a global function with example usage: get_prng(seed) and returns: np