class_docstr: Perform weighted sum to generate output features according to scores.
  Modified from `PAConv <https://github.com/CVMI-Lab/PAConv/tree/main/
  scene_seg/lib/paconv_lib/src/gpu>`_.
  
  This is a memory-efficient CUDA implementation of assign_scores operation,
  which first transform all point features with weight bank, then assemble
  neighbor features with ``knn_idx`` and perform weighted sum of ``scores``.
  
  See the `paper <https://arxiv.org/pdf/2103.14635.pdf>`_ appendix Sec. D for
      more detailed descriptions.
  
  Note:
      This implementation assumes using ``neighbor`` kernel input, which is
          (point_features - center_features, point_features).
      See https://github.com/CVMI-Lab/PAConv/blob/main/scene_seg/model/
      pointnet2/paconv.py#L128 for more details.
class_docstr_tok: ["Perform","weighted","sum","to","generate","output","features","according","to","scores",".","Modified","from","`","PAConv","<","https",":","\/\/github.com\/CVMI-Lab\/PAConv\/tree\/main\/","scene_seg\/lib\/paconv_lib\/src\/gpu",">","`","_",".","This","is","a","memory-efficient","CUDA","implementation","of","assign_scores","operation",",","which","first","transform","all","point","features","with","weight","bank",",","then","assemble","neighbor","features","with","``","knn_idx","``","and","perform","weighted","sum","of","``","scores","``",".","See","the","`","paper","<","https",":","\/\/arxiv.org\/pdf\/2103.14635.pdf",">","`","_","appendix","Sec",".","D","for","more","detailed","descriptions",".","Note",":","This","implementation","assumes","using","``","neighbor","``","kernel","input",",","which","is","(","point_features","-","center_features",",","point_features",")",".","See","https",":","\/\/github.com\/CVMI-Lab\/PAConv\/blob\/main\/scene_seg\/model\/","pointnet2\/paconv.py","#","L128","for","more","details","."]
docstr: Args:
      scores (torch.Tensor): (B, npoint, K, M), predicted scores to
          aggregate weight matrices in the weight bank.
          ``npoint`` is the number of sampled centers.
          ``K`` is the number of queried neighbors.
          ``M`` is the number of weight matrices in the weight bank.
      point_features (torch.Tensor): (B, N, M, out_dim)
          Pre-computed point features to be aggregated.
      center_features (torch.Tensor): (B, N, M, out_dim)
          Pre-computed center features to be aggregated.
      knn_idx (torch.Tensor): (B, npoint, K), index of sampled kNN.
          We assume the first idx in each row is the idx of the center.
      aggregate (str, optional): Aggregation method.
          Can be 'sum', 'avg' or 'max'. Defaults: 'sum'.
  
  Returns:
      torch.Tensor: (B, out_dim, npoint, K), the aggregated features.
docstr_tok: ["Args",":","scores","(","torch.Tensor",")",":","(","B",",","npoint",",","K",",","M",")",",","predicted","scores","to","aggregate","weight","matrices","in","the","weight","bank",".","``","npoint","``","is","the","number","of","sampled","centers",".","``","K","``","is","the","number","of","queried","neighbors",".","``","M","``","is","the","number","of","weight","matrices","in","the","weight","bank",".","point_features","(","torch.Tensor",")",":","(","B",",","N",",","M",",","out_dim",")","Pre-computed","point","features","to","be","aggregated",".","center_features","(","torch.Tensor",")",":","(","B",",","N",",","M",",","out_dim",")","Pre-computed","center","features","to","be","aggregated",".","knn_idx","(","torch.Tensor",")",":","(","B",",","npoint",",","K",")",",","index","of","sampled","kNN",".","We","assume","the","first","idx","in","each","row","is","the","idx","of","the","center",".","aggregate","(","str",",","optional",")",":","Aggregation","method",".","Can","be","'sum","'",",","'avg","'","or","'max","'",".","Defaults",":","'sum","'",".","Returns",":","torch.Tensor",":","(","B",",","out_dim",",","npoint",",","K",")",",","the","aggregated","features","."]
code: def forward(
      ctx,
      scores,
      point_features,
      center_features,
      knn_idx,
      aggregate="sum",
  ):
      """
      Args:
          scores (torch.Tensor): (B, npoint, K, M), predicted scores to
              aggregate weight matrices in the weight bank.
              ``npoint`` is the number of sampled centers.
              ``K`` is the number of queried neighbors.
              ``M`` is the number of weight matrices in the weight bank.
          point_features (torch.Tensor): (B, N, M, out_dim)
              Pre-computed point features to be aggregated.
          center_features (torch.Tensor): (B, N, M, out_dim)
              Pre-computed center features to be aggregated.
          knn_idx (torch.Tensor): (B, npoint, K), index of sampled kNN.
              We assume the first idx in each row is the idx of the center.
          aggregate (str, optional): Aggregation method.
              Can be 'sum', 'avg' or 'max'. Defaults: 'sum'.
  
      Returns:
          torch.Tensor: (B, out_dim, npoint, K), the aggregated features.
      """
      agg = {"sum": 0, "avg": 1, "max": 2}
  
      B, N, M, out_dim = point_features.size()
      _, npoint, K, _ = scores.size()
  
      output = point_features.new_zeros((B, out_dim, npoint, K))
      ext_module.assign_score_withk_forward(
          point_features.contiguous(),
          center_features.contiguous(),
          scores.contiguous(),
          knn_idx.contiguous(),
          output,
          B=B,
          N0=N,
          N1=npoint,
          M=M,
          K=K,
          O=out_dim,
          aggregate=agg[aggregate],
      )
  
      ctx.save_for_backward(
          output, point_features, center_features, scores, knn_idx
      )
      ctx.agg = agg[aggregate]
  
      return output
  
imports: from torch.autograd import Function
  from ..utils import ext_loader
label_desc: Use image node_id 1 for calling the AssignScoreWithK obj's underlying member method code with example usage: obj.forward(ctx, scores, point_features, center_features, knn_idx, aggregate) and returns: output