code: def define_prng_impl(
      *,
      key_shape: Shape,
      seed: Callable[[Array], Array],
      split: Callable[[Array, Shape], Array],
      random_bits: Callable[[Array, int, Shape], Array],
      fold_in: Callable[[Array, int], Array],
      name: str = "<unnamed>",
      tag: str = "?",
  ) -> Hashable:
      return random.PRNGSpec(
          prng.PRNGImpl(
              key_shape,
              seed,
              split,
              random_bits,
              fold_in,
              name=name,
              tag=tag,
          )
      )
  
imports: from typing import Callable
  from collections.abc import Hashable
  from jax import Array
  from jax._src import prng
  from jax._src import random
label_desc: Use image node_id 1 for calling a global function with example usage: define_prng_impl() and returns: random