code: def sample_inverse_transform_stratified_1d(
      ray_org: TensorLike,
      ray_dir: TensorLike,
      z_values_init: TensorLike,
      weights_init: TensorLike,
      n_samples: int,
      combine_z_values=True,
      name: str = "sample_inverse_transform_stratified_1d",
  ):
      """Sample points on a ray using inverse transform stratified sampling.
  
      The rays are defined by their origin and direction. Along each ray, there are
      M samples (provided as 1D distances from the ray origin) and the corresponding
      weights (probabilities) that facilitate the inverse transform sampling.
  
      Args:
        ray_org: A tensor of shape `[A1, ..., An, 3]`,
          where the last dimension represents the 3D position of the ray origin.
        ray_dir: A tensor of shape `[A1, ..., An, 3]`,
          where the last dimension represents the 3D direction of the ray.
        z_values_init: A tensor of shape `[A1, ..., An, M]`,
          where the last dimension is the location of M points along the ray.
        weights_init: A tensor of shape `[A1, ..., An, M]`,
          where the last dimension is the density of M points along the ray.
        n_samples: A number M to sample on the ray.
        combine_z_values: Wether to combine the new 1D samples with
          the initial points.
        name: A name for this op that defaults to "stratified_sampling".
  
      Returns:
        A tensor of shape `[A1, ..., An, M, 3]` indicating the M points on the ray
          and a tensor of shape `[A1, ..., An, M]` for the Z values on the points.
      """
      with tf.name_scope(name):
          shape.check_static(
              tensor=ray_org,
              tensor_name="ray_org",
              has_dim_equals=(-1, 3),
          )
          shape.check_static(
              tensor=ray_dir,
              tensor_name="ray_dir",
              has_dim_equals=(-1, 3),
          )
          shape.compare_batch_dimensions(
              tensors=(ray_org, ray_dir, z_values_init, weights_init),
              tensor_names=(
                  "ray_org",
                  "ray_dir",
                  "z_values_init",
                  "weights_init",
              ),
              last_axes=-2,
              broadcast_compatible=False,
          )
          shape.compare_dimensions(
              tensors=(z_values_init, weights_init),
              tensor_names=("z_values_init", "weights_init"),
              axes=-1,
          )
  
          bin_start = z_values_init[..., :-1]
          bin_width = z_values_init[..., 1:] - z_values_init[..., :-1]
          bin_weights = 0.5 * (
              weights_init[..., 1:] + weights_init[..., :-1]
          )
          random_z_values = sampling.inverse_transform_stratified_1d(
              bin_start, bin_width, bin_weights, n_samples
          )
          random_z_values = tf.stop_gradient(random_z_values)
          if combine_z_values:
              z_values_final = tf.sort(
                  tf.concat([z_values_init, random_z_values], -1), -1
              )
          else:
              z_values_final = tf.sort(random_z_values, -1)
          points3d = _points_from_z_values(
              ray_org, ray_dir, z_values_final
          )
          return points3d, z_values_final
  
file_path: graphics/tensorflow_graphics/geometry/representation/ray.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Tuple, Union
  from six.moves import range
  import tensorflow
  from tensorflow_graphics.math import sampling
  from tensorflow_graphics.math import vector
  from tensorflow_graphics.util import asserts
  from tensorflow_graphics.util import export_api
  from tensorflow_graphics.util import shape
  from tensorflow_graphics.util.type_alias import TensorLike
label_desc: Use image node_id 4 for calling a global function with example usage: sample_inverse_transform_stratified_1d(ray_org, ray_dir, z_values_init, weights_init, n_samples, combine_z_values, name) and returns: points3d,  z_values_final