code: def volumetric_rendering(
      rgb,
      sigma,
      z_vals,
      dirs,
      use_white_background,
      sample_at_infinity=True,
      return_weights=False,
      eps=1e-10,
  ):
      """Volumetric Rendering Function.
  
      Args:
        rgb: an array of size (B,S,3) containing the RGB color values.
        sigma: an array of size (B,S,1) containing the densities.
        z_vals: an array of size (B,S) containing the z-coordinate of the samples.
        dirs: an array of size (B,3) containing the directions of rays.
        use_white_background: whether to assume a white background or not.
        sample_at_infinity: if True adds a sample at infinity.
        return_weights: if True returns the weights in the dictionary.
        eps: a small number to prevent numerical issues.
  
      Returns:
        A dictionary containing:
          rgb: an array of size (B,3) containing the rendered colors.
          depth: an array of size (B,) containing the rendered depth.
          acc: an array of size (B,) containing the accumulated density.
          weights: an array of size (B,S) containing the weight of each sample.
      """
      # TODO(keunhong): remove this hack.
      last_sample_z = 1e10 if sample_at_infinity else 1e-19
      dists = jnp.concatenate(
          [
              z_vals[..., 1:] - z_vals[..., :-1],
              jnp.broadcast_to([last_sample_z], z_vals[..., :1].shape),
          ],
          -1,
      )
      dists = dists * jnp.linalg.norm(dirs[..., None, :], axis=-1)
      alpha = 1.0 - jnp.exp(-sigma * dists)
      # Prepend a 1.0 to make this an 'exclusive' cumprod as in `tf.math.cumprod`.
      accum_prod = jnp.concatenate(
          [
              jnp.ones_like(alpha[..., :1], alpha.dtype),
              jnp.cumprod(1.0 - alpha[..., :-1] + eps, axis=-1),
          ],
          axis=-1,
      )
      weights = alpha * accum_prod
  
      rgb = (weights[..., None] * rgb).sum(axis=-2)
      exp_depth = (weights * z_vals).sum(axis=-1)
      med_depth = compute_depth_map(weights, z_vals)
      acc = weights.sum(axis=-1)
      if use_white_background:
          rgb = rgb + (1.0 - acc[..., None])
  
      if sample_at_infinity:
          acc = weights[..., :-1].sum(axis=-1)
  
      out = {
          "rgb": rgb,
          "depth": exp_depth,
          "med_depth": med_depth,
          "acc": acc,
      }
      if return_weights:
          out["weights"] = weights
      return out
  
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 2 for calling a global function with example usage: volumetric_rendering(rgb, sigma, z_vals, dirs, use_white_background, sample_at_infinity, return_weights, eps) and returns: out