code: def _xy_to_ray_bundle(
      cameras: CamerasBase,
      xy_grid: torch.Tensor,
      min_depth: float,
      max_depth: float,
      n_pts_per_ray: int,
      unit_directions: bool,
      stratified_sampling: bool = False,
  ) -> RayBundle:
      """
      Extends the `xy_grid` input of shape `(batch_size, ..., 2)` to rays.
      This adds to each xy location in the grid a vector of `n_pts_per_ray` depths
      uniformly spaced between `min_depth` and `max_depth`.
  
      The extended grid is then unprojected with `cameras` to yield
      ray origins, directions and depths.
  
      Args:
          cameras: cameras object representing a batch of cameras.
          xy_grid: torch.tensor grid of image xy coords.
          min_depth: The minimum depth of each ray-point.
          max_depth: The maximum depth of each ray-point.
          n_pts_per_ray: The number of points sampled along each ray.
          unit_directions: whether to normalize direction vectors in ray bundle.
          stratified_sampling: if True, performs stratified sampling in n_pts_per_ray
              bins for each ray; otherwise takes n_pts_per_ray deterministic points
              on each ray with uniform offsets.
      """
      batch_size = xy_grid.shape[0]
      spatial_size = xy_grid.shape[1:-1]
      n_rays_per_image = spatial_size.numel()
  
      # ray z-coords
      rays_zs = xy_grid.new_empty((0,))
      if n_pts_per_ray > 0:
          depths = torch.linspace(
              min_depth,
              max_depth,
              n_pts_per_ray,
              dtype=xy_grid.dtype,
              device=xy_grid.device,
          )
          rays_zs = depths[None, None].expand(
              batch_size, n_rays_per_image, n_pts_per_ray
          )
  
          if stratified_sampling:
              rays_zs = _jiggle_within_stratas(rays_zs)
  
      # make two sets of points at a constant depth=1 and 2
      to_unproject = torch.cat(
          (
              xy_grid.view(batch_size, 1, n_rays_per_image, 2)
              .expand(batch_size, 2, n_rays_per_image, 2)
              .reshape(batch_size, n_rays_per_image * 2, 2),
              torch.cat(
                  (
                      xy_grid.new_ones(batch_size, n_rays_per_image, 1),
                      2.0
                      * xy_grid.new_ones(
                          batch_size, n_rays_per_image, 1
                      ),
                  ),
                  dim=1,
              ),
          ),
          dim=-1,
      )
  
      # unproject the points
      unprojected = cameras.unproject_points(
          to_unproject, from_ndc=True
      )
  
      # split the two planes back
      rays_plane_1_world = unprojected[:, :n_rays_per_image]
      rays_plane_2_world = unprojected[:, n_rays_per_image:]
  
      # directions are the differences between the two planes of points
      rays_directions_world = rays_plane_2_world - rays_plane_1_world
  
      # origins are given by subtracting the ray directions from the first plane
      rays_origins_world = rays_plane_1_world - rays_directions_world
  
      if unit_directions:
          rays_directions_world = F.normalize(
              rays_directions_world, dim=-1
          )
  
      return RayBundle(
          rays_origins_world.view(batch_size, *spatial_size, 3),
          rays_directions_world.view(batch_size, *spatial_size, 3),
          rays_zs.view(batch_size, *spatial_size, n_pts_per_ray),
          xy_grid,
      )
  
imports: import warnings
  from typing import Optional, Tuple, Union
  import torch
  from pytorch3d.common.compat import meshgrid_ij
  from pytorch3d.ops import padded_to_packed
  from pytorch3d.renderer.cameras import CamerasBase
  from pytorch3d.renderer.implicit.utils import HeterogeneousRayBundle, RayBundle
  from torch.nn import functional
label_desc: Use image node_id 4 for calling a global function with example usage: _xy_to_ray_bundle(cameras, xy_grid, min_depth, max_depth, n_pts_per_ray, unit_directions, stratified_sampling) and returns: RayBundle