docstr: march rays to generate points (forward only)
  Args:
      rays_o/d: float, [N, 3]
      bound: float, scalar
      density_bitfield: uint8: [CHHH // 8]
      C: int
      H: int
      nears/fars: float, [N]
      step_counter: int32, (2), used to count the actual number of generated points.
      mean_count: int32, estimated mean steps to accelerate training. (but will randomly drop rays if the actual point count exceeded this threshold.)
      perturb: bool
      align: int, pad output so its size is dividable by align, set to -1 to disable.
      force_all_rays: bool, ignore step_counter and mean_count, always calculate all rays. Useful if rendering the whole image, instead of some rays.
      dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
      max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
  Returns:
      xyzs: float, [M, 3], all generated points' coords. (all rays concated, need to use `rays` to extract points belonging to each ray)
      dirs: float, [M, 3], all generated points' view dirs.
      ts: float, [M, 2], all generated points' ts.
      rays: int32, [N, 2], all rays' (point_offset, point_count), e.g., xyzs[rays[i, 0]:(rays[i, 0] + rays[i, 1])] --> points belonging to rays[i, 0]
docstr_tok: ["march","rays","to","generate","points","(","forward","only",")","Args",":","rays_o\/d",":","float",",","[","N",",","3","]","bound",":","float",",","scalar","density_bitfield",":","uint8",":","[","CHHH","\/\/","8","]","C",":","int","H",":","int","nears\/fars",":","float",",","[","N","]","step_counter",":","int32",",","(","2",")",",","used","to","count","the","actual","number","of","generated","points",".","mean_count",":","int32",",","estimated","mean","steps","to","accelerate","training",".","(","but","will","randomly","drop","rays","if","the","actual","point","count","exceeded","this","threshold",".",")","perturb",":","bool","align",":","int",",","pad","output","so","its","size","is","dividable","by","align",",","set","to","-1","to","disable",".","force_all_rays",":","bool",",","ignore","step_counter","and","mean_count",",","always","calculate","all","rays",".","Useful","if","rendering","the","whole","image",",","instead","of","some","rays",".","dt_gamma",":","float",",","called","cone_angle","in","instant-ngp",",","exponentially","accelerate","ray","marching","if",">","0",".","(","very","significant","effect",",","but","generally","lead","to","worse","performance",")","max_steps",":","int",",","max","number","of","sampled","points","along","each","ray",",","also","affect","min_stepsize",".","Returns",":","xyzs",":","float",",","[","M",",","3","]",",","all","generated","points","'","coords",".","(","all","rays","concated",",","need","to","use","`","rays","`","to","extract","points","belonging","to","each","ray",")","dirs",":","float",",","[","M",",","3","]",",","all","generated","points","'","view","dirs",".","ts",":","float",",","[","M",",","2","]",",","all","generated","points","'","ts",".","rays",":","int32",",","[","N",",","2","]",",","all","rays","'","(","point_offset",",","point_count",")",",","e.g.",",","xyzs","[","rays","[","i",",","0","]",":","(","rays","[","i",",","0","]","+","rays","[","i",",","1","]",")","]","--",">","points","belonging","to","rays","[","i",",","0","]"]
code: def forward(
      ctx,
      rays_o,
      rays_d,
      bound,
      density_bitfield,
      C,
      H,
      nears,
      fars,
      perturb=False,
      dt_gamma=0,
      max_steps=1024,
      contract=False,
  ):
      """march rays to generate points (forward only)
      Args:
          rays_o/d: float, [N, 3]
          bound: float, scalar
          density_bitfield: uint8: [CHHH // 8]
          C: int
          H: int
          nears/fars: float, [N]
          step_counter: int32, (2), used to count the actual number of generated points.
          mean_count: int32, estimated mean steps to accelerate training. (but will randomly drop rays if the actual point count exceeded this threshold.)
          perturb: bool
          align: int, pad output so its size is dividable by align, set to -1 to disable.
          force_all_rays: bool, ignore step_counter and mean_count, always calculate all rays. Useful if rendering the whole image, instead of some rays.
          dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
          max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
      Returns:
          xyzs: float, [M, 3], all generated points' coords. (all rays concated, need to use `rays` to extract points belonging to each ray)
          dirs: float, [M, 3], all generated points' view dirs.
          ts: float, [M, 2], all generated points' ts.
          rays: int32, [N, 2], all rays' (point_offset, point_count), e.g., xyzs[rays[i, 0]:(rays[i, 0] + rays[i, 1])] --> points belonging to rays[i, 0]
      """
  
      if not rays_o.is_cuda:
          rays_o = rays_o.cuda()
      if not rays_d.is_cuda:
          rays_d = rays_d.cuda()
      if not density_bitfield.is_cuda:
          density_bitfield = density_bitfield.cuda()
  
      rays_o = rays_o.float().contiguous().view(-1, 3)
      rays_d = rays_d.float().contiguous().view(-1, 3)
      density_bitfield = density_bitfield.contiguous()
  
      N = rays_o.shape[0]  # num rays
  
      step_counter = torch.zeros(
          1, dtype=torch.int32, device=rays_o.device
      )  # point counter, ray counter
  
      if perturb:
          noises = torch.rand(
              N, dtype=rays_o.dtype, device=rays_o.device
          )
      else:
          noises = torch.zeros(
              N, dtype=rays_o.dtype, device=rays_o.device
          )
  
      # first pass: write rays, get total number of points M to render
      rays = torch.empty(
          N, 2, dtype=torch.int32, device=rays_o.device
      )  # id, offset, num_steps
      get_backend().march_rays_train(
          rays_o,
          rays_d,
          density_bitfield,
          bound,
          contract,
          dt_gamma,
          max_steps,
          N,
          C,
          H,
          nears,
          fars,
          None,
          None,
          None,
          rays,
          step_counter,
          noises,
      )
  
      # allocate based on M
      M = step_counter.item()
      # print(M, N)
      # print(rays[:, 0].max())
  
      xyzs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
      dirs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
      ts = torch.zeros(M, 2, dtype=rays_o.dtype, device=rays_o.device)
  
      # second pass: write outputs
      get_backend().march_rays_train(
          rays_o,
          rays_d,
          density_bitfield,
          bound,
          contract,
          dt_gamma,
          max_steps,
          N,
          C,
          H,
          nears,
          fars,
          xyzs,
          dirs,
          ts,
          rays,
          step_counter,
          noises,
      )
  
      return xyzs, dirs, ts, rays
  
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 _march_rays_train obj's underlying member method code with example usage: obj.forward(ctx, rays_o, rays_d, bound, density_bitfield, C, H, nears, fars, perturb, dt_gamma, max_steps, contract) and returns: xyzs,  dirs,  ts,  rays