def _uniform_unit_norm(dimension, shape, dtype, seed):
  """Returns a batch of points chosen uniformly from the unit hypersphere."""
  # This works because the Gaussian distribution is spherically symmetric.
  # raw shape: shape + [dimension]
  raw = normal.Normal(
      loc=dtype_util.as_numpy_dtype(dtype)(0),
      scale=dtype_util.as_numpy_dtype(dtype)(1)).sample(
          tf.concat([shape, [dimension]], axis=0), seed=seed())
  unit_norm = raw / tf.norm(tensor=raw, ord=2, axis=-1)[..., tf.newaxis]
  return unit_norm
Returns a batch of points chosen uniformly from the unit hypersphere.