def threefry_2x32_prf(key, x: jnp.ndarray) -> jnp.ndarray:
  """Apply the threefry PRF to an array of inputs.

  This function is vectorized over x.
  For threefry_2x32: K = X = uint32[2]

  Args:
    key: uint32[2] the key of the PRF
    x: uint32[..., 2] the inputs

  Returns:
    y: uint32[..., 2] the outputs
  """
  if not (key.shape == (2,) and key.dtype == jnp.uint32):
    raise TypeError('key must be uint32[2]', key)
  if not (x.shape[-1:] == (2,) and x.dtype == jnp.uint32):
    raise TypeError('x must be uint32[..., 2]', x)
  # Threefry-2x32 expects this weird format:
  x_3f = jnp.moveaxis(x, source=-1, destination=0).flatten()
  y_3f = jex.random.threefry_2x32(key, x_3f)
  y = jnp.moveaxis(
      jnp.reshape(y_3f, (2,) + x.shape[:-1]), source=0, destination=-1)
  return y
Apply the threefry PRF to an array of inputs.

  This function is vectorized over x.
  For threefry_2x32: K = X = uint32[2]

  Args:
    key: uint32[2] the key of the PRF
    x: uint32[..., 2] the inputs

  Returns:
    y: uint32[..., 2] the outputs
  
trax/layers/research/position_encodings.py