code: def threefry_2x32_prange(key, lo: int = 0, hi: int = 2):
      """Splits a key into a stream of random keys.
  
      This uses the little-endian counter mode.
  
      Args:
        key: uint32[2] the key to split
        lo: the range to start extracting from
        hi: the range to stop extracting from
  
      Returns:
        keys: uint32[hi - lo, 2] the split keys
      """
      if not (key.shape == (2,) and key.dtype == jnp.uint32):
          raise ValueError("key must be uint32[2]")
      if not hi < 2**32:
          # You shouldn't really be using more than half the key size anyways.
          raise NotImplementedError("only 32-bit sizes are supported")
      # Create a 64-bit counter:
      i_lo = jnp.arange(lo, hi, dtype=jnp.uint32)
      i_hi = jnp.zeros_like(i_lo)
      i = jnp.stack([i_lo, i_hi], axis=-1)
      return threefry_2x32_prf(key, i)
  
imports: import logging
  import jax
  import jax.extend
  import numpy
  import trax
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 2 for calling a global function with example usage: threefry_2x32_prange(key, lo, hi) and returns: threefry_2x32_prf