docstr: flatten rays
  Args:
      rays: [N, 2], all rays' (point_offset, point_count),
      M: scalar, int, count of points (we cannot get this info from rays unfortunately...)
  Returns:
      res: [M], flattened ray index.
docstr_tok: ["flatten","rays","Args",":","rays",":","[","N",",","2","]",",","all","rays","'","(","point_offset",",","point_count",")",",","M",":","scalar",",","int",",","count","of","points","(","we","can","not","get","this","info","from","rays","unfortunately","...",")","Returns",":","res",":","[","M","]",",","flattened","ray","index","."]
code: def forward(ctx, rays, M):
      """flatten rays
      Args:
          rays: [N, 2], all rays' (point_offset, point_count),
          M: scalar, int, count of points (we cannot get this info from rays unfortunately...)
      Returns:
          res: [M], flattened ray index.
      """
      if not rays.is_cuda:
          rays = rays.cuda()
      rays = rays.contiguous()
  
      N = rays.shape[0]
  
      res = torch.zeros(M, dtype=torch.int, device=rays.device)
  
      get_backend().flatten_rays(rays, N, M, res)
  
      return res
  
imports: import torch
  from torch.autograd import Function
  from torch.cuda.amp import custom_bwd, custom_fwd
label_desc: Use image node_id 1 for calling the _flatten_rays obj's underlying member method code with example usage: obj.forward(ctx, rays, M) and returns: res