def normal(key, shape, dtype=tf.float32):
  """Sample standard-normal random values.

  Args:
    key: the RNG key.
    shape: the shape of the result.
    dtype: the dtype of the result.

  Returns:
    Random values in standard-normal distribution.
  """
  key = tf_np.asarray(key, dtype=_RNG_KEY_DTYPE)
  return tf_np.asarray(
      tf.random.stateless_normal(shape, seed=_key2seed(key), dtype=dtype))
Sample standard-normal random values.

  Args:
    key: the RNG key.
    shape: the shape of the result.
    dtype: the dtype of the result.

  Returns:
    Random values in standard-normal distribution.
  
trax/tf_numpy/extensions/extensions.py