def get_tensors(self, device): white_values = torch.from_numpy(np.ctypeslib.as_array(self.white_values, shape=(self.num_active_white_features,))).pin_memory().to(device=device, non_blocking=True) black_values = torch.from_numpy(np.ctypeslib.as_array(self.black_values, shape=(self.num_active_black_features,))).pin_memory().to(device=device, non_blocking=True) iw = torch.transpose(torch.from_numpy(np.ctypeslib.as_array(self.white, shape=(self.num_active_white_features, 2))).pin_memory().to(device=device, non_blocking=True), 0, 1).long() ib = torch.transpose(torch.from_numpy(np.ctypeslib.as_array(self.black, shape=(self.num_active_white_features, 2))).pin_memory().to(device=device, non_blocking=True), 0, 1).long() us = torch.from_numpy(np.ctypeslib.as_array(self.is_white, shape=(self.size, 1))).pin_memory().to(device=device, non_blocking=True) them = 1.0 - us outcome = torch.from_numpy(np.ctypeslib.as_array(self.outcome, shape=(self.size, 1))).pin_memory().to(device=device, non_blocking=True) score = torch.from_numpy(np.ctypeslib.as_array(self.score, shape=(self.size, 1))).pin_memory().to(device=device, non_blocking=True) white = torch._sparse_coo_tensor_unsafe(iw, white_values, (self.size, self.num_inputs)) black = torch._sparse_coo_tensor_unsafe(ib, black_values, (self.size, self.num_inputs)) white._coalesced_(True) black._coalesced_(True) return us, them, white, black, outcome, score
def sparse_matrix(data, index, shape, force_format=False): fmt = index[0] if fmt != 'coo': raise TypeError('Pytorch backend only supports COO format. But got %s.' % fmt) # NOTE: use _sparse_coo_tensor_unsafe to avoid unnecessary boundary check spmat = th._sparse_coo_tensor_unsafe(index[1], data, shape) # No conversion is required. return spmat, None
def _rebuild_sparse_tensor(layout, data): if layout == torch.sparse_coo: indices, values, size = data result = torch._sparse_coo_tensor_unsafe(indices, values, size) _sparse_tensors_to_validate.append(result) return result raise NotImplementedError("rebuilding sparse tensor for layout %s" % (layout))
def test_factory_override(self): class A(TorchFunctionMode): def __torch_function__(self, *args, **kwargs): return -1 with torch.overrides.push_torch_function_mode(A): self.assertEqual(torch.tensor([1]), -1) self.assertEqual(torch.sparse_coo_tensor(1, 1, 1), -1) self.assertEqual(torch.sparse_csr_tensor(1, 1, 1), -1) self.assertEqual(torch._sparse_coo_tensor_unsafe(1, 1, (1, 1)), -1) self.assertEqual(torch._sparse_csr_tensor_unsafe(1, 1, 1, (1, 1)), -1) self.assertEqual(torch.as_tensor([1]), -1)
def _rebuild_sparse_tensor(layout, data): """ Rebuilds a sparse tensor from its sparse storage representation. Args: layout (str): The sparse storage layout of the tensor. data (tuple): The tensor's sparse storage representation. """ if layout == torch.sparse_coo: indices, values, size = data result = torch._sparse_coo_tensor_unsafe(indices, values, size) _sparse_tensors_to_validate.append(result) return result raise NotImplementedError("rebuilding sparse tensor for layout %s" % (layout))