code: def noise_regularize(key, raw, noise_std, use_stratified_sampling):
      """Regularize the density prediction by adding gaussian noise.
  
      Args:
        key: jnp.ndarray(float32), [2,], random number generator.
        raw: jnp.ndarray(float32), [batch_size, num_coarse_samples, 4].
        noise_std: float, std dev of noise added to regularize sigma output.
        use_stratified_sampling: add noise only if use_stratified_sampling is True.
  
      Returns:
        raw: jnp.ndarray(float32), [batch_size, num_coarse_samples, 4], updated raw.
      """
      if (
          (noise_std is not None)
          and noise_std > 0.0
          and use_stratified_sampling
      ):
          unused_key, key = random.split(key)
          noise = (
              random.normal(key, raw[..., 3:4].shape, dtype=raw.dtype)
              * noise_std
          )
          raw = jnp.concatenate(
              [raw[..., :3], raw[..., 3:4] + noise], axis=-1
          )
      return raw
  
file_path: nerfies/nerfies/model_utils.py
imports: from flax import linen
  from flax import optim
  from flax import struct
  from jax import lax
  from jax import random
  import jax.numpy
label_desc: Use image node_id 8 for calling a global function with example usage: noise_regularize(key, raw, noise_std, use_stratified_sampling) and returns: raw