docs: def sample_truncated_gaussian(mu=0, sigma=1, lb=-np.Inf, ub=np.Inf):
      """
      Sample a truncated normal with the specified params. This
      is not the most stable way but it works as long as the
      truncation region is not too far from the mean.
      """
      # Broadcast arrays to be of the same shape
      mu, sigma, lb, ub = np.broadcast_arrays(mu, sigma, lb, ub)
      shp = mu.shape
      if np.allclose(sigma, 0.0):
          return mu
  
      cdflb = normal_cdf(lb, mu, sigma)
      cdfub = normal_cdf(ub, mu, sigma)
  
      # Sample uniformly from the CDF
      cdfsamples = cdflb + np.random.rand(*shp) * (cdfub-cdflb)
  
      # Clip the CDF samples so that we can invert them
      cdfsamples = np.clip(cdfsamples, 1e-15, 1-1e-15)
      zs = -np.sqrt(2) * special.erfcinv(2 * cdfsamples)
  
      # Transform the standard normal samples
      xs = sigma * zs + mu
      xs = np.clip(xs, lb, ub)
  
      return xs