def METHOD_NAME(ctx,
            scores: torch.Tensor,
            point_features: torch.Tensor,
            center_features: torch.Tensor,
            knn_idx: torch.Tensor,
            aggregate: str = 'sum') -> torch.Tensor:
    """
    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