code: def piecewise_constant_pdf(
      key, bins, weights, num_coarse_samples, use_stratified_sampling
  ):
      """Piecewise-Constant PDF sampling.
  
      Args:
        key: jnp.ndarray(float32), [2,], random number generator.
        bins: jnp.ndarray(float32), [batch_size, n_bins + 1].
        weights: jnp.ndarray(float32), [batch_size, n_bins].
        num_coarse_samples: int, the number of samples.
        use_stratified_sampling: bool, use use_stratified_sampling samples.
  
      Returns:
        z_samples: jnp.ndarray(float32), [batch_size, num_coarse_samples].
      """
      eps = 1e-5
  
      # Get pdf
      weights += eps  # prevent nans
      pdf = weights / weights.sum(axis=-1, keepdims=True)
      cdf = jnp.cumsum(pdf, axis=-1)
      cdf = jnp.concatenate(
          [jnp.zeros(list(cdf.shape[:-1]) + [1]), cdf], axis=-1
      )
  
      # Take uniform samples
      if use_stratified_sampling:
          u = random.uniform(
              key, list(cdf.shape[:-1]) + [num_coarse_samples]
          )
      else:
          u = jnp.linspace(0.0, 1.0, num_coarse_samples)
          u = jnp.broadcast_to(
              u, list(cdf.shape[:-1]) + [num_coarse_samples]
          )
  
      # Invert CDF. This takes advantage of the fact that `bins` is sorted.
      mask = u[..., None, :] >= cdf[..., :, None]
  
      def minmax(x):
          x0 = jnp.max(
              jnp.where(mask, x[..., None], x[..., :1, None]), -2
          )
          x1 = jnp.min(
              jnp.where(~mask, x[..., None], x[..., -1:, None]), -2
          )
          x0 = jnp.minimum(x0, x[..., -2:-1])
          x1 = jnp.maximum(x1, x[..., 1:2])
          return x0, x1
  
      bins_g0, bins_g1 = minmax(bins)
      cdf_g0, cdf_g1 = minmax(cdf)
  
      denom = cdf_g1 - cdf_g0
      denom = jnp.where(denom < eps, 1.0, denom)
      t = (u - cdf_g0) / denom
      z_samples = bins_g0 + t * (bins_g1 - bins_g0)
  
      # Prevent gradient from backprop-ing through samples
      return lax.stop_gradient(z_samples)
  
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 3 for calling a global function with example usage: piecewise_constant_pdf(key, bins, weights, num_coarse_samples, use_stratified_sampling) and returns: lax, x0,  x1