def METHOD_NAME(ctx, grad_out):
    """
    Args:
        grad_out (torch.Tensor): (B, out_dim, npoint, K)
    Returns:
        grad_scores (torch.Tensor): (B, npoint, K, M)
        grad_point_features (torch.Tensor): (B, N, M, out_dim)
        grad_center_features (torch.Tensor): (B, N, M, out_dim)
    """
    _, point_features, center_features, scores, knn_idx = ctx.saved_tensors
    agg = ctx.agg
    B, N, M, out_dim = point_features.size()
    _, npoint, K, _ = scores.size()
    grad_point_features = point_features.new_zeros(point_features.shape)
    grad_center_features = center_features.new_zeros(center_features.shape)
    grad_scores = scores.new_zeros(scores.shape)
    ext_module.assign_score_withk_backward(
        grad_out.contiguous(),
        point_features.contiguous(),
        center_features.contiguous(),
        scores.contiguous(),
        knn_idx.contiguous(),
        grad_point_features,
        grad_center_features,
        grad_scores,
        B=B,
        N0=N,
        N1=npoint,
        M=M,
        K=K,
        O=out_dim,
        aggregate=agg)
    return grad_scores, grad_point_features, \
        grad_center_features, None, None