Class MinkowskiSPMMFunction

MinkowskiEngine/sparse_matrix_functions.py:124–169

class MinkowskiSPMMFunction(Function):
    @staticmethod
    def forward(
        ctx,
        rows: torch.Tensor,
        cols: torch.Tensor,
        vals: torch.Tensor,
        size: torch.Size,
        mat: torch.Tensor,
        cuda_spmm_alg: int = 1,
    ):
        ctx.misc_args = size, cuda_spmm_alg
        ctx.save_for_backward(rows, cols, vals)
        result = spmm(
            rows,
            cols,
            vals,
            size,
            mat,
            is_sorted=False,
            cuda_spmm_alg=cuda_spmm_alg,
        )
        return result

    @staticmethod
    def backward(ctx, grad: torch.Tensor):
        size, cuda_spmm_alg = ctx.misc_args
        rows, cols, vals = ctx.saved_tensors
        new_size = torch.Size([size[1], size[0]])
        grad = spmm(
            cols,
            rows,
            vals,
            new_size,
            grad,
            is_sorted=False,
            cuda_spmm_alg=cuda_spmm_alg,
        )
        return (
            None,
            None,
            None,
            None,
            grad,
            None,
        )


class MinkowskiSPMMAverageFunction(Function):