def sample_along_rays(key, origins, directions, radii, num_samples, near, far,
                      randomized, lindisp, ray_shape):
  """Stratified sampling along the rays.

  Args:
    key: jnp.ndarray, random generator key.
    origins: jnp.ndarray(float32), [batch_size, 3], ray origins.
    directions: jnp.ndarray(float32), [batch_size, 3], ray directions.
    radii: jnp.ndarray(float32), [batch_size, 3], ray radii.
    num_samples: int.
    near: jnp.ndarray, [batch_size, 1], near clip.
    far: jnp.ndarray, [batch_size, 1], far clip.
    randomized: bool, use randomized stratified sampling.
    lindisp: bool, sampling linearly in disparity rather than depth.
    ray_shape: string, which shape ray to assume.

  Returns:
    t_vals: jnp.ndarray, [batch_size, num_samples], sampled z values.
    means: jnp.ndarray, [batch_size, num_samples, 3], sampled means.
    covs: jnp.ndarray, [batch_size, num_samples, 3, 3], sampled covariances.
  """
  batch_size = origins.shape[0]

  t_vals = jnp.linspace(0., 1., num_samples + 1)
  if lindisp:
    t_vals = 1. / (1. / near * (1. - t_vals) + 1. / far * t_vals)
  else:
    t_vals = near * (1. - t_vals) + far * t_vals

  if randomized:
    mids = 0.5 * (t_vals[..., 1:] + t_vals[..., :-1])
    upper = jnp.concatenate([mids, t_vals[..., -1:]], -1)
    lower = jnp.concatenate([t_vals[..., :1], mids], -1)
    t_rand = random.uniform(key, [batch_size, num_samples + 1])
    t_vals = lower + (upper - lower) * t_rand
  else:
    # Broadcast t_vals to make the returned shape consistent.
    t_vals = jnp.broadcast_to(t_vals, [batch_size, num_samples + 1])
  means, covs = cast_rays(t_vals, origins, directions, radii, ray_shape)
  return t_vals, (means, covs)
Stratified sampling along the rays.

  Args:
    key: jnp.ndarray, random generator key.
    origins: jnp.ndarray(float32), [batch_size, 3], ray origins.
    directions: jnp.ndarray(float32), [batch_size, 3], ray directions.
    radii: jnp.ndarray(float32), [batch_size, 3], ray radii.
    num_samples: int.
    near: jnp.ndarray, [batch_size, 1], near clip.
    far: jnp.ndarray, [batch_size, 1], far clip.
    randomized: bool, use randomized stratified sampling.
    lindisp: bool, sampling linearly in disparity rather than depth.
    ray_shape: string, which shape ray to assume.

  Returns:
    t_vals: jnp.ndarray, [batch_size, num_samples], sampled z values.
    means: jnp.ndarray, [batch_size, num_samples, 3], sampled means.
    covs: jnp.ndarray, [batch_size, num_samples, 3, 3], sampled covariances.
  
internal/mip.py