code: def sample_along_rays(
      key,
      origins,
      directions,
      num_coarse_samples,
      near,
      far,
      use_stratified_sampling,
      use_linear_disparity,
  ):
      """Stratified sampling along the rays.
  
      Args:
        key: jnp.ndarray, random generator key.
        origins: ray origins.
        directions: ray directions.
        num_coarse_samples: int.
        near: float, near clip.
        far: float, far clip.
        use_stratified_sampling: use stratified sampling.
        use_linear_disparity: sampling linearly in disparity rather than depth.
  
      Returns:
        z_vals: jnp.ndarray, [batch_size, num_coarse_samples], sampled z values.
        points: jnp.ndarray, [batch_size, num_coarse_samples, 3], sampled points.
      """
      batch_size = origins.shape[0]
  
      t_vals = jnp.linspace(0.0, 1.0, num_coarse_samples)
      if not use_linear_disparity:
          z_vals = near * (1.0 - t_vals) + far * t_vals
      else:
          z_vals = 1.0 / (
              1.0 / near * (1.0 - t_vals) + 1.0 / far * t_vals
          )
      if use_stratified_sampling:
          mids = 0.5 * (z_vals[..., 1:] + z_vals[..., :-1])
          upper = jnp.concatenate([mids, z_vals[..., -1:]], -1)
          lower = jnp.concatenate([z_vals[..., :1], mids], -1)
          t_rand = random.uniform(key, [batch_size, num_coarse_samples])
          z_vals = lower + (upper - lower) * t_rand
      else:
          # Broadcast z_vals to make the returned shape consistent.
          z_vals = jnp.broadcast_to(
              z_vals[None, ...], [batch_size, num_coarse_samples]
          )
  
      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 1 for calling a global function with example usage: sample_along_rays(key, origins, directions, num_coarse_samples, near, far, use_stratified_sampling, use_linear_disparity) and returns: z_vals,  unknown