So far I wrote:
# works correctly only for non-negative tensors
def save_topk(x, k, dim = -1, indices_dtype = None, values_dtype = None):
    topk = x.topk(k, dim = dim)
    return dict(k = k, dim = dim, shape = x.shape, indices = topk.indices.to(dtype = indices_dtype), values = topk.values.to(dtype = values_dtype), dtype = x.dtype)

def load_topk(saved, **kwargs):
    return torch.zeros(saved['shape'], dtype = saved['dtype'], **kwargs).scatter_(saved['dim'], saved['indices'].long(), saved['values'].to(dtype = saved['dtype']))

x = torch.rand(3, 4, 5)
y = save_topk(x, 2, dim = -1, indices_dtype = torch.uint8, values_dtype = torch.bfloat16)
z = load_topk(y)

I guess at this point my question is: can stock hybrid sparse-dense tensors help eliminate custom tensor saving format (while preserving uint8/bfloat16 custom dtypes) and loading logic?