def resample_along_rays(key, origins, directions, radii, t_vals, weights,
                        randomized, ray_shape, stop_grad, resample_padding):
  """Resampling.

  Args:
    key: jnp.ndarray(float32), [2,], random number generator.
    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.
    t_vals: jnp.ndarray(float32), [batch_size, num_samples+1].
    weights: jnp.array(float32), weights for t_vals
    randomized: bool, use randomized samples.
    ray_shape: string, which kind of shape to assume for the ray.
    stop_grad: bool, whether or not to backprop through sampling.
    resample_padding: float, added to the weights before normalizing.

  Returns:
    t_vals: jnp.ndarray(float32), [batch_size, num_samples+1].
    points: jnp.ndarray(float32), [batch_size, num_samples, 3].
  """
  # Do a blurpool.
  weights_pad = jnp.concatenate([
      weights[..., :1],
      weights,
      weights[..., -1:],
  ],
                                axis=-1)
  weights_max = jnp.maximum(weights_pad[..., :-1], weights_pad[..., 1:])
  weights_blur = 0.5 * (weights_max[..., :-1] + weights_max[..., 1:])

  # Add in a constant (the sampling function will renormalize the PDF).
  weights = weights_blur + resample_padding

  new_t_vals = math.sorted_piecewise_constant_pdf(
      key,
      t_vals,
      weights,
      t_vals.shape[-1],
      randomized,
  )
  if stop_grad:
    new_t_vals = lax.stop_gradient(new_t_vals)
  means, covs = cast_rays(new_t_vals, origins, directions, radii, ray_shape)
  return new_t_vals, (means, covs)
Resampling.

  Args:
    key: jnp.ndarray(float32), [2,], random number generator.
    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.
    t_vals: jnp.ndarray(float32), [batch_size, num_samples+1].
    weights: jnp.array(float32), weights for t_vals
    randomized: bool, use randomized samples.
    ray_shape: string, which kind of shape to assume for the ray.
    stop_grad: bool, whether or not to backprop through sampling.
    resample_padding: float, added to the weights before normalizing.

  Returns:
    t_vals: jnp.ndarray(float32), [batch_size, num_samples+1].
    points: jnp.ndarray(float32), [batch_size, num_samples, 3].
  
internal/mip.py