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)