code: def sample_pdf(
      key,
      bins,
      weights,
      origins,
      directions,
      z_vals,
      num_coarse_samples,
      use_stratified_sampling,
  ):
      """Hierarchical 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].
        origins: ray origins.
        directions: ray directions.
        z_vals: jnp.ndarray(float32), [batch_size, n_coarse_samples].
        num_coarse_samples: int, the number of samples.
        use_stratified_sampling: bool, use use_stratified_sampling samples.
  
      Returns:
        z_vals: jnp.ndarray(float32),
          [batch_size, n_coarse_samples + num_fine_samples].
        points: jnp.ndarray(float32),
          [batch_size, n_coarse_samples + num_fine_samples, 3].
      """
      z_samples = piecewise_constant_pdf(
          key,
          bins,
          weights,
          num_coarse_samples,
          use_stratified_sampling,
      )
      # Compute united z_vals and sample points
      z_vals = jnp.sort(
          jnp.concatenate([z_vals, z_samples], axis=-1), axis=-1
      )
      return z_vals, (
          origins[..., None, :]
          + z_vals[..., None] * directions[..., None, :]
      )
  
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 4 for calling a global function with example usage: sample_pdf(key, bins, weights, origins, directions, z_vals, num_coarse_samples, use_stratified_sampling) and returns: z_vals,  unknown