I don’t know to be honest. There seems to be a promising package here: https://github.com/cornellius-gp/gpytorch 13
Check out BDSMM on line 83:

  
      github.com
  
  
    cornellius-gp/gpytorch/blob/master/gpytorch/utils/sparse.py 5
#!/usr/bin/env python3

from operator import mul

import torch


def make_sparse_from_indices_and_values(interp_indices, interp_values, num_rows):
    """
    This produces a sparse tensor with a fixed number of non-zero entries in each column.

    Args:
        - interp_indices - Tensor (batch_size) x num_cols x n_nonzero_entries
            A matrix which has the indices of the nonzero_entries for each column
        - interp_values - Tensor (batch_size) x num_cols x n_nonzero_entries
            The corresponding values
        - num_rows - the number of rows in the result matrix

    Returns:
        - SparseTensor - (batch_size) x num_cols x num_rows


  This file has been truncated. show original

  
  
    
    
  
  


Seems like they change the indices of what would be the 3d sparse matrix to a 2d by adding the number of rows of the each matrix to the next set of indices.