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))
