def reduce(x: torch.Tensor):
if x.is_sparse:
assert x.ndim == 2, "Only 2-dimensional tensors are supported"
mode = get_mode(x)
return rebuild_dispatcher, args_dispatcher(mode, x)
else:
return x.__reduce_ex__(pickle.DEFAULT_PROTOCOL)
