code: def get_rays(
      poses: torch.Tensor,
      intrinsics: torch.Tensor,
      height: int,
      width: int,
      num_samples: Optional[int] = None,
      error_map: Optional[torch.Tensor] = None,
      device: torch.device = "cuda",
  ) -> Dict[str, torch.Tensor]:
      """
      Generates rays from camera poses and intrinsics.
  
      Args:
          poses (torch.Tensor): Camera poses, shape [B, 4, 4] (cam2world).
          intrinsics (torch.Tensor): Intrinsic camera parameters [fx, fy, cx, cy].
          height (int): Height of the image.
          width (int): Width of the image.
          num_samples: Number of rays to sample, default is None for all rays.
          error_map: Optional tensor to use for non-uniform sampling of rays.
          device (torch.device): Device on which to generate the rays.
  
      Returns:
          Dict[str, torch.Tensor]: A dictionary containing the following keys:
              - 'rays_o': Origin of the rays, shape [B, N, 3]
              - 'rays_d': Directions of the rays, shape [B, N, 3]
              - 'inds': Indices of the rays, shape [B, N] (if N > 0)
              - 'inds_coarse': Coarse indices of the rays, shape [B, N] (if error_map is not None)
      """
  
      batch_size = poses.shape[0]
      fx, fy, cx, cy = intrinsics
  
      i, j = torch.meshgrid(
          torch.linspace(0, width - 1, width, device=device),
          torch.linspace(0, height - 1, height, device=device),
          indexing="ij",
      )
      i = (
          i.t()
          .reshape([1, height * width])
          .expand([batch_size, height * width])
          + 0.5
      )
      j = (
          j.t()
          .reshape([1, height * width])
          .expand([batch_size, height * width])
          + 0.5
      )
  
      results = {}
  
      if num_samples is not None:
          num_samples = min(num_samples, height * width)
  
          if error_map is None:
              sampled_indices = torch.randint(
                  0, height * width, size=[num_samples], device=device
              )
              sampled_indices = sampled_indices.expand(
                  [batch_size, num_samples]
              )
          else:
              (
                  sampled_indices,
                  sampled_indices_coarse,
              ) = non_uniform_sampling(
                  error_map=error_map,
                  num_samples=num_samples,
                  height=height,
                  width=width,
                  device=device,
              )
              results["sampled_indices_coarse"] = sampled_indices_coarse
  
          i = torch.gather(i, -1, sampled_indices)
          j = torch.gather(j, -1, sampled_indices)
          results["sampled_indices"] = sampled_indices
      else:
          sampled_indices = torch.arange(
              height * width, device=device
          ).expand([batch_size, height * width])
  
      zs = torch.full_like(i, -1.0)
      xs = -(i - cx) / fx * zs
      ys = (j - cy) / fy * zs
      directions = torch.stack((xs, ys, zs), dim=-1)
  
      rays_d = directions @ poses[:, :3, :3].transpose(-1, -2)
      rays_o = poses[..., :3, 3].unsqueeze(-2).expand_as(rays_d)
  
      rays_o = rays_o.view(-1, height, width, 3)
      rays_d = rays_d.view(-1, height, width, 3)
  
      return rays_o, rays_d
  
file_path: NeMo/nemo/collections/multimodal/data/nerf/utils.py
imports: from typing import Dict, Optional
  import numpy
  import torch
  import torch.nn.functional
label_desc: Use image node_id 4 for calling a global function with example usage: get_rays(poses, intrinsics, height, width, num_samples, error_map, device) and returns: rays_o,  rays_d