def binary_cross_entropy(input, target, weight=None, size_average=True):
n, c, h, w = input.size()
nt, ht, wt = target.size()
# Handle inconsistent size between input and target
if h > ht and w > wt: # upsample labels
target = target.unsequeeze(1)
target = F.upsample(target, size=(h, w), mode="nearest")
target = target.sequeeze(1)
elif h < ht and w < wt: # upsample images
input = F.upsample(input, size=(ht, wt), mode="bilinear")
elif h != ht and w != wt:
raise Exception("Only support upsampling")
input = input.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)
target = target.view(-1).long()
target_flip = torch.bitwise_xor(target, 1)
target = torch.cat((target.unsqueeze_(-1), target_flip.unsqueeze_(-1)),
dim=1)
print(input.size())
print(target.size())
loss = nn.BCEWithLogitsLoss()
return loss(input, target)
def __init__(self, config, attention_type):
super().__init__()
max_positions = config.max_position_embeddings
bias = torch.tril(
torch.ones((max_positions, max_positions),
dtype=torch.uint8)).view(1, 1, max_positions,
max_positions)
# local causal self attention is a sliding window where each token can only attend to the previous
# window_size tokens. This is implemented by updating the causal mask such that for each token
# all other tokens are masked except the previous window_size tokens.
if attention_type == "local":
bias = torch.bitwise_xor(bias, torch.tril(bias,
-config.window_size))
self.register_buffer("bias", bias)
self.register_buffer("masked_bias", torch.tensor(-1e9))
self.attn_dropout = nn.Dropout(config.attention_dropout)
self.resid_dropout = nn.Dropout(config.resid_dropout)
self.embed_dim = config.hidden_size
self.num_heads = config.num_heads
self.head_dim = self.embed_dim // self.num_heads
if self.head_dim * self.num_heads != self.embed_dim:
raise ValueError(
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
f" {self.num_heads}).")
self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
def __init__(self, config, attention_type):
super().__init__()
max_positions = config.n_pos
bias = torch.tril(
torch.ones((max_positions, max_positions),
dtype=torch.uint8)).view(1, 1, max_positions,
max_positions)
if attention_type == "local":
bias = torch.bitwise_xor(bias, torch.tril(bias, -config.s_win))
self.register_buffer("bias", bias)
self.register_buffer("masked_bias", torch.tensor(-1e9))
self.attn_dropout = qc.Dropout(config.drop_attn)
self.drop_resid = qc.Dropout(config.drop_resid)
self.embed_dim = config.d_model
self.n_heads = config.n_heads
self.head_dim = self.embed_dim // self.n_heads
if self.head_dim * self.n_heads != self.embed_dim:
raise ValueError(
f"embed_dim must be divisible by n_heads (got `embed_dim`: {self.embed_dim} and `n_heads`: {self.n_heads})."
)
self.k_proj = qc.Linear(self.embed_dim, self.embed_dim, bias=False)
self.v_proj = qc.Linear(self.embed_dim, self.embed_dim, bias=False)
self.q_proj = qc.Linear(self.embed_dim, self.embed_dim, bias=False)
self.out_proj = qc.Linear(self.embed_dim, self.embed_dim, bias=True)
def bitwise_or(a, b):
"""
Bitwise OR of Tensors via int intermediate
"""
a = a.add(.5).mul(256.).round().int()
b = b.add(.5).mul(256.).round().int()
# Bitwise OR on integers, convert back to float
return torch.bitwise_xor(a, b).div(256.).sub(.5)
def forward(ctx, input_raw, xormask):
input_int = torch.round(input_raw).to(
dtype=torch.uint8) # round != cast to int
xormask_int = xormask.to(dtype=torch.uint8)
ctx.save_for_backward(input_int, xormask_int)
res = torch.bitwise_xor(input_int,
xormask_int).to(dtype=input_raw.dtype)
return res
def bitwise_xor(input_, other):
"""Wrapper of `torch.bitwise_xor`.
Parameters
----------
input_ : DTensor
The first operand.
other : DTensor
The second operand.
"""
return torch.bitwise_xor(input_._data, other._data)
def fnv_hash_vec(arr):
"""
FNV64-1A
"""
assert arr.ndim == 3
# Floor first for negative coordinates
arr = arr.copy()
arr = arr.astype(np.uint64, copy=False)
hashed_arr = np.uint64(14695981039346656037) * torch.ones((arr.size(0), arr.size(1)), dtype=np.uint64)
for j in range(arr.shape[1]): # loop on each coord channel
hashed_arr *= np.uint64(1099511628211)
hashed_arr = torch.bitwise_xor(hashed_arr, arr[:, j])
return hashed_arr
def perturb(ten: torch.tensor, repr_width, p):
# TODO smart kernel for GPU/CPU vs multiplication to concatenate mask
# bitrank vector
# compact_sample = torch.sum(compact_sample * bitrank, dim=-1)
ten_repr = Cpp_Pert.generateTensorMask(ten, repr_width, p)
print(ten_repr)
ten_np = ten_repr.numpy()
packed = np.packbits(
ten_np.astype(int), axis=-1,
bitorder="little") # Packing bits in order to create the mask
ten_packed = torch.from_numpy(packed)
mask = torch.flatten(
ten_packed,
start_dim=-2) # Removing the extra dimension caused by the bit packing
ten.data = torch.bitwise_xor(ten, mask).data
def cnn_kernel_3x3(img, kernel, N):
"""Unipolar kernel convolve -> AND gates
Bipolar kernel convolve -> XNOR gates"""
is_bp = np.any(kernel < 0)
h, w = img.shape
rng = bs.SC_RNG()
rng_type = rng.bs_bp_uniform if is_bp else rng.bs_uniform
img_bs = img_io.img_to_bs(img, rng_type, bs_len=N) #Correlated at +1
img_bs = torch.from_numpy(img_bs).to(device)
rng.reset()
kernel_bs = img_io.img_to_bs(kernel, rng_type, bs_len=N, scale=False)
kernel_bs = torch.from_numpy(kernel_bs).to(device)
nb = N>>3
rc_mat = torch.cuda.ByteTensor(h-2, w-2, nb).fill_(0)
m = torch.cuda.ByteTensor(3, 3, nb).fill_(0)
z = torch.cuda.ByteTensor(nb).fill_(0)
for i in range(h-2):
for j in range(w-2):
for k in range(3):
for l in range(3):
if is_bp:
m[k][l] = torch.bitwise_not(torch.bitwise_xor(img_bs[i + k][j + l], kernel_bs[k][l]))
else:
m[k][l] = torch.bitwise_and(img_bs[i + k][j + l], kernel_bs[k][l])
#mux sum tree
l1_1 = mux_p_cuda(m[0][0], m[0][1], 0.5)
l1_2 = mux_p_cuda(m[0][2], m[1][0], 0.5)
l1_3 = mux_p_cuda(m[1][1], m[1][2], 0.5)
l1_4 = mux_p_cuda(m[2][0], m[2][1], 0.5)
l2_1 = mux_p_cuda(l1_1, l1_2, 0.5)
l2_2 = mux_p_cuda(l1_3, l1_4, 0.5)
l3 = mux_p_cuda(l2_1, l2_2, 0.5)
rc_mat[i][j] = mux_p_cuda(l3, mux_p_cuda(m[2][2], z, 1.0/8.0), 1.0/9.0)
#mean_type = bs.bs_mean_bp if is_bp else bs.bs_mean
#rc_mat = rc_mat.to(cpu)
#img_io.disp_img(img_io.bs_to_img(rc_mat, mean_type, scaling=9))
return bs.get_corr_mat_cuda(rc_mat.view((h-2)*(w-2), nb)).to(cpu).numpy()
def robert_cross_maj_cuda(rands, x11, x12, x21, x22):
xor1 = torch.bitwise_xor(x11, x22)
xor2 = torch.bitwise_xor(x12, x21)
return maj_p_cuda(xor1, xor2, rands)
def __threefry64(X_0: torch.Tensor, X_1: torch.Tensor, seed: int):
"""
Counter-based pseudo random number generator. Based on a 12-round Threefry "encryption" algorithm [1]. This is the
64-bit version.
Parameters
----------
X_0 : torch.Tensor
Upper bits of the to be encoded random sequence
X_1 : torch.Tensor
Lower bits of the to be encoded random sequence
Returns
-------
random_numbers : tuple(torch.Tensor (int64))
Two vectors with num_samples / 2 (rounded-up) pseudo random numbers.
References
----------
[1] Salmon, John K., Moraes, Mark A., Dror, Ron O. and Shaw, David E., "Parallel random numbers: as easy as 1, 2, 3"
Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis,
p. 16, 2011
"""
samples = len(X_0)
# set up key buffer
ks_0 = torch.full((samples, ), seed, dtype=torch.int64, device=X_0.device)
ks_1 = torch.full((samples, ), seed, dtype=torch.int64, device=X_1.device)
ks_2 = torch.full((samples, ),
2004413935125273122,
dtype=torch.int64,
device=X_0.device)
# ks_2 ^= ks_0
# ks_2 ^= ks_1
ks_2 = torch.bitwise_xor(torch.bitwise_xor(ks_2, ks_0), ks_1)
# initialize output using the key
X_0 += ks_0
X_1 += ks_1
# perform rounds
# round 1
X_0 += X_1
X_1 = (X_1 << 16) | ((X_1 >> 48) & 0xFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 2
X_0 += X_1
X_1 = (X_1 << 42) | ((X_1 >> 22) & 0x3FFFFFFFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 3
X_0 += X_1
X_1 = (X_1 << 12) | ((X_1 >> 52) & 0xFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 4
X_0 += X_1
X_1 = (X_1 << 31) | ((X_1 >> 33) & 0x7FFFFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# inject key
X_0 += ks_1
X_1 += ks_2 + 1
# round 5
X_0 += X_1
X_1 = (X_1 << 16) | ((X_1 >> 48) & 0xFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 6
X_0 += X_1
X_1 = (X_1 << 32) | ((X_1 >> 32) & 0xFFFFFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 7
X_0 += X_1
X_1 = (X_1 << 24) | ((X_1 >> 40) & 0xFFFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# round 8
X_0 += X_1
X_1 = (X_1 << 21) | ((X_1 >> 43) & 0x1FFFFF)
# X_1 ^= X_0
X_1 = torch.bitwise_xor(X_1, X_0)
# inject key
# X_0 += ks_2; X_1 += (ks_0 + 2)
#
# X_0 += X_1; X_1 = (X_1 << 16) | (X_1 >> 48); X_1 ^= X_0 # round 9
# X_0 += X_1; X_1 = (X_1 << 42) | (X_1 >> 22); X_1 ^= X_0 # round 10
# X_0 += X_1; X_1 = (X_1 << 12) | (X_1 >> 52); X_1 ^= X_0 # round 11
# X_0 += X_1; X_1 = (X_1 << 31) | (X_1 >> 33); X_1 ^= X_0 # round 12
# inject key
X_0 += ks_0
X_1 += ks_1 + 3
return X_0, X_1
print(torch.clamp(a, -0.5, 0.5))
# trigonometric functions and their inverses
angles = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
sines = torch.sin(angles)
inverses = torch.asin(sines)
print('\nSine and arcsine:')
print(angles)
print(sines)
print(inverses)
# bitwise operations
print('\nBitwise XOR:')
b = torch.tensor([1, 5, 11])
c = torch.tensor([2, 7, 10])
print(torch.bitwise_xor(b, c))
# comparisons:
print('\nBroadcasted, element-wise equality comparison:')
d = torch.tensor([[1., 2.], [3., 4.]])
e = torch.ones(1, 2) # many comparison ops support broadcasting!
print(torch.eq(d, e)) # returns a tensor of type bool
# reductions:
print('\nReduction ops:')
print(torch.max(d)) # returns a single-element tensor
print(torch.max(d).item()) # extracts the value from the returned tensor
print(torch.mean(d)) # average
print(torch.std(d)) # standard deviation
print(torch.prod(d)) # product of all numbers
print(torch.unique(torch.tensor([1, 2, 1, 2, 1, 2]))) # filter unique elements
def perturb(self, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
return torch.bitwise_xor(tensor, mask)
def forward(self, a, b):
c = torch.bitwise_xor(a, b)
d = torch.bitwise_xor(a, c)
return d
def precompute_unique_patches_all_styles(config):
if torch.cuda.is_available():
device = torch.device('cuda')
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
result_dir = "output_for_eval"
if not os.path.exists(result_dir):
print("ERROR: result_dir does not exist! " + result_dir)
exit(-1)
patches_dir = "unique_patches"
if not os.path.exists(patches_dir):
os.makedirs(patches_dir)
#load style shapes
fin = open("splits/" + config.data_style + ".txt")
styleset_names = [name.strip() for name in fin.readlines()]
fin.close()
styleset_len = len(styleset_names)
print("precompute_unique_patches:", config.data_style, styleset_len)
buffer_size = 256 * 256 * 16 #change the buffer size if the input voxel is large
patch_size = 12
if not config.asymmetry:
padding_size = 8 - patch_size // 2
else:
padding_size = 0
#prepare dictionary for style shapes
dict_style_patches = []
dict_style_patches_edge = []
dict_style_patches_dilated = []
dict_style_patches_tensor = []
patch_num_total = 0
for style_id in range(styleset_len):
data_dict = h5py.File(
patches_dir + "/style_" + str(style_id) + ".hdf5", 'r')
patches = data_dict['patches'][:]
patches_edge = data_dict['patches_edge'][:]
patches_dilated = data_dict['patches_dilated'][:]
data_dict.close()
patch_num = len(patches)
patch_num_total += patch_num
dict_style_patches.append(patches)
dict_style_patches_edge.append(patches_edge)
dict_style_patches_dilated.append(patches_dilated)
patches_tensor = torch.from_numpy(patches).to(device).view(
patch_num, -1).bool()
dict_style_patches_tensor.append(patches_tensor)
pointer_total = 0
dict_style_patches_len = []
for style_id in range(styleset_len):
print(style_id, styleset_len)
this_patches = dict_style_patches[style_id]
this_patches_edge = dict_style_patches_edge[style_id]
this_patches_dilated = dict_style_patches_dilated[style_id]
this_patches_tensor = dict_style_patches_tensor[style_id]
patch_num = len(this_patches)
pointer = 1
for patch_id in range(1, patch_num):
notduplicated_flag = True
for compare_id in range(0, style_id + 1):
compare_patches_tensor = dict_style_patches_tensor[compare_id]
if compare_id != style_id:
notduplicated = torch.bitwise_xor(
this_patches_tensor[patch_id:patch_id + 1],
compare_patches_tensor[:dict_style_patches_len[
compare_id]]).any(1).all().item()
else:
notduplicated = torch.bitwise_xor(
this_patches_tensor[patch_id:patch_id + 1],
compare_patches_tensor[:pointer]).any(1).all().item()
if not notduplicated:
notduplicated_flag = False
break
if notduplicated_flag:
this_patches[pointer] = this_patches[patch_id]
this_patches_edge[pointer] = this_patches_edge[patch_id]
this_patches_dilated[pointer] = this_patches_dilated[patch_id]
this_patches_tensor.data[pointer] = this_patches_tensor.data[
patch_id]
pointer += 1
pointer_total += pointer
dict_style_patches_len.append(pointer)
print("before --", patch_num, patch_num_total)
print("after --", pointer, pointer_total)
hdf5_file = h5py.File(
patches_dir + "/global_unique_" + config.data_style + ".hdf5", 'w')
for style_id in range(styleset_len):
pointer = dict_style_patches_len[style_id]
hdf5_file.create_dataset("patches_" + str(style_id),
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file.create_dataset("patches_edge_" + str(style_id),
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file.create_dataset("patches_dilated_" + str(style_id),
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file["patches_" +
str(style_id)][:] = dict_style_patches[style_id][:pointer]
hdf5_file[
"patches_edge_" +
str(style_id)][:] = dict_style_patches_edge[style_id][:pointer]
hdf5_file[
"patches_dilated_" +
str(style_id)][:] = dict_style_patches_dilated[style_id][:pointer]
hdf5_file.close()
def precompute_unique_patches_per_style(config):
if torch.cuda.is_available():
device = torch.device('cuda')
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
result_dir = "output_for_eval"
if not os.path.exists(result_dir):
print("ERROR: result_dir does not exist! " + result_dir)
exit(-1)
patches_dir = "unique_patches"
if not os.path.exists(patches_dir):
os.makedirs(patches_dir)
#load style shapes
fin = open("splits/" + config.data_style + ".txt")
styleset_names = [name.strip() for name in fin.readlines()]
fin.close()
styleset_len = len(styleset_names)
print("precompute_unique_patches:", config.data_style, styleset_len)
buffer_size = 256 * 256 * 16 #change the buffer size if the input voxel is large
patch_size = 12
if not config.asymmetry:
padding_size = 8 - patch_size // 2
else:
padding_size = 0
#prepare dictionary for style shapes
for style_id in range(styleset_len):
print(style_id, styleset_len)
voxel_model_file = open(
result_dir + "/style_" + str(style_id) + ".binvox", 'rb')
style_shape = binvox_rw.read_as_3d_array(voxel_model_file,
fix_coords=False).data.astype(
np.uint8)
style_shape = np.ascontiguousarray(style_shape[:, :, padding_size:])
style_shape_tensor = torch.from_numpy(style_shape).to(
device).unsqueeze(0).unsqueeze(0).float()
style_shape_edge_tensor = F.max_pool3d(
-style_shape_tensor, kernel_size=3, stride=1,
padding=1) + style_shape_tensor
style_shape_dilated_tensor = F.max_pool3d(style_shape_edge_tensor,
kernel_size=3,
stride=1,
padding=1)
style_shape_edge = style_shape_edge_tensor.detach().cpu().numpy()[0, 0]
style_shape_edge = np.round(style_shape_edge).astype(np.uint8)
style_shape_dilated = style_shape_dilated_tensor.detach().cpu().numpy(
)[0, 0]
style_shape_dilated = np.round(style_shape_dilated).astype(np.uint8)
patches = np.zeros([buffer_size, patch_size, patch_size, patch_size],
np.uint8)
patches_edge = np.zeros(
[buffer_size, patch_size, patch_size, patch_size], np.uint8)
patches_dilated = np.zeros(
[buffer_size, patch_size, patch_size, patch_size], np.uint8)
patch_num = cutils.get_patches_edge_dilated(
style_shape, style_shape_edge, style_shape_dilated, patches,
patches_edge, patches_dilated, patch_size)
patches = np.copy(patches[:patch_num])
patches_edge = np.copy(patches_edge[:patch_num])
patches_dilated = np.copy(patches_dilated[:patch_num])
patches_tensor = torch.from_numpy(patches).to(device).view(
patch_num, -1).bool()
pointer = 0
for patch_id in range(patch_num - 1):
notduplicated = torch.bitwise_xor(
patches_tensor[patch_id:patch_id + 1],
patches_tensor[patch_id + 1:]).any(1).all().item()
if notduplicated:
patches[pointer] = patches[patch_id]
patches_edge[pointer] = patches_edge[patch_id]
patches_dilated[pointer] = patches_dilated[patch_id]
pointer += 1
patch_id = patch_num - 1
patches[pointer] = patches[patch_id]
patches_edge[pointer] = patches_edge[patch_id]
patches_dilated[pointer] = patches_dilated[patch_id]
pointer += 1
print("before --", patch_num)
print("after --", pointer)
hdf5_file = h5py.File(
patches_dir + "/style_" + str(style_id) + ".hdf5", 'w')
hdf5_file.create_dataset("patches",
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file.create_dataset("patches_edge",
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file.create_dataset("patches_dilated",
[pointer, patch_size, patch_size, patch_size],
np.uint8,
compression=9)
hdf5_file["patches"][:] = patches[:pointer]
hdf5_file["patches_edge"][:] = patches_edge[:pointer]
hdf5_file["patches_dilated"][:] = patches_dilated[:pointer]
hdf5_file.close()
def pointwise_ops(self):
a = torch.randn(4)
b = torch.randn(4)
t = torch.tensor([-1, -2, 3], dtype=torch.int8)
r = torch.tensor([0, 1, 10, 0], dtype=torch.int8)
t = torch.tensor([-1, -2, 3], dtype=torch.int8)
s = torch.tensor([4, 0, 1, 0], dtype=torch.int8)
f = torch.zeros(3)
g = torch.tensor([-1, 0, 1])
w = torch.tensor([0.3810, 1.2774, -0.2972, -0.3719, 0.4637])
return (
torch.abs(torch.tensor([-1, -2, 3])),
torch.absolute(torch.tensor([-1, -2, 3])),
torch.acos(a),
torch.arccos(a),
torch.acosh(a.uniform_(1.0, 2.0)),
torch.add(a, 20),
torch.add(a, torch.randn(4, 1), alpha=10),
torch.addcdiv(torch.randn(1, 3),
torch.randn(3, 1),
torch.randn(1, 3),
value=0.1),
torch.addcmul(torch.randn(1, 3),
torch.randn(3, 1),
torch.randn(1, 3),
value=0.1),
torch.angle(a),
torch.asin(a),
torch.arcsin(a),
torch.asinh(a),
torch.arcsinh(a),
torch.atan(a),
torch.arctan(a),
torch.atanh(a.uniform_(-1.0, 1.0)),
torch.arctanh(a.uniform_(-1.0, 1.0)),
torch.atan2(a, a),
torch.bitwise_not(t),
torch.bitwise_and(t, torch.tensor([1, 0, 3], dtype=torch.int8)),
torch.bitwise_or(t, torch.tensor([1, 0, 3], dtype=torch.int8)),
torch.bitwise_xor(t, torch.tensor([1, 0, 3], dtype=torch.int8)),
torch.ceil(a),
torch.clamp(a, min=-0.5, max=0.5),
torch.clamp(a, min=0.5),
torch.clamp(a, max=0.5),
torch.clip(a, min=-0.5, max=0.5),
torch.conj(a),
torch.copysign(a, 1),
torch.copysign(a, b),
torch.cos(a),
torch.cosh(a),
torch.deg2rad(
torch.tensor([[180.0, -180.0], [360.0, -360.0], [90.0,
-90.0]])),
torch.div(a, b),
torch.divide(a, b, rounding_mode="trunc"),
torch.divide(a, b, rounding_mode="floor"),
torch.digamma(torch.tensor([1.0, 0.5])),
torch.erf(torch.tensor([0.0, -1.0, 10.0])),
torch.erfc(torch.tensor([0.0, -1.0, 10.0])),
torch.erfinv(torch.tensor([0.0, 0.5, -1.0])),
torch.exp(torch.tensor([0.0, math.log(2.0)])),
torch.exp2(torch.tensor([0.0, math.log(2.0), 3.0, 4.0])),
torch.expm1(torch.tensor([0.0, math.log(2.0)])),
torch.fake_quantize_per_channel_affine(
torch.randn(2, 2, 2),
(torch.randn(2) + 1) * 0.05,
torch.zeros(2),
1,
0,
255,
),
torch.fake_quantize_per_tensor_affine(a, 0.1, 0, 0, 255),
torch.float_power(torch.randint(10, (4, )), 2),
torch.float_power(torch.arange(1, 5), torch.tensor([2, -3, 4,
-5])),
torch.floor(a),
# torch.floor_divide(torch.tensor([4.0, 3.0]), torch.tensor([2.0, 2.0])),
# torch.floor_divide(torch.tensor([4.0, 3.0]), 1.4),
torch.fmod(torch.tensor([-3, -2, -1, 1, 2, 3]), 2),
torch.fmod(torch.tensor([1, 2, 3, 4, 5]), 1.5),
torch.frac(torch.tensor([1.0, 2.5, -3.2])),
torch.randn(4, dtype=torch.cfloat).imag,
torch.ldexp(torch.tensor([1.0]), torch.tensor([1])),
torch.ldexp(torch.tensor([1.0]), torch.tensor([1, 2, 3, 4])),
torch.lerp(torch.arange(1.0, 5.0),
torch.empty(4).fill_(10), 0.5),
torch.lerp(
torch.arange(1.0, 5.0),
torch.empty(4).fill_(10),
torch.full_like(torch.arange(1.0, 5.0), 0.5),
),
torch.lgamma(torch.arange(0.5, 2, 0.5)),
torch.log(torch.arange(5) + 10),
torch.log10(torch.rand(5)),
torch.log1p(torch.randn(5)),
torch.log2(torch.rand(5)),
torch.logaddexp(torch.tensor([-1.0]), torch.tensor([-1, -2, -3])),
torch.logaddexp(torch.tensor([-100.0, -200.0, -300.0]),
torch.tensor([-1, -2, -3])),
torch.logaddexp(torch.tensor([1.0, 2000.0, 30000.0]),
torch.tensor([-1, -2, -3])),
torch.logaddexp2(torch.tensor([-1.0]), torch.tensor([-1, -2, -3])),
torch.logaddexp2(torch.tensor([-100.0, -200.0, -300.0]),
torch.tensor([-1, -2, -3])),
torch.logaddexp2(torch.tensor([1.0, 2000.0, 30000.0]),
torch.tensor([-1, -2, -3])),
torch.logical_and(r, s),
torch.logical_and(r.double(), s.double()),
torch.logical_and(r.double(), s),
torch.logical_and(r, s, out=torch.empty(4, dtype=torch.bool)),
torch.logical_not(torch.tensor([0, 1, -10], dtype=torch.int8)),
torch.logical_not(
torch.tensor([0.0, 1.5, -10.0], dtype=torch.double)),
torch.logical_not(
torch.tensor([0.0, 1.0, -10.0], dtype=torch.double),
out=torch.empty(3, dtype=torch.int16),
),
torch.logical_or(r, s),
torch.logical_or(r.double(), s.double()),
torch.logical_or(r.double(), s),
torch.logical_or(r, s, out=torch.empty(4, dtype=torch.bool)),
torch.logical_xor(r, s),
torch.logical_xor(r.double(), s.double()),
torch.logical_xor(r.double(), s),
torch.logical_xor(r, s, out=torch.empty(4, dtype=torch.bool)),
torch.logit(torch.rand(5), eps=1e-6),
torch.hypot(torch.tensor([4.0]), torch.tensor([3.0, 4.0, 5.0])),
torch.i0(torch.arange(5, dtype=torch.float32)),
torch.igamma(a, b),
torch.igammac(a, b),
torch.mul(torch.randn(3), 100),
torch.multiply(torch.randn(4, 1), torch.randn(1, 4)),
torch.mvlgamma(torch.empty(2, 3).uniform_(1.0, 2.0), 2),
torch.tensor([float("nan"),
float("inf"), -float("inf"), 3.14]),
torch.nan_to_num(w),
torch.nan_to_num(w, nan=2.0),
torch.nan_to_num(w, nan=2.0, posinf=1.0),
torch.neg(torch.randn(5)),
# torch.nextafter(torch.tensor([1, 2]), torch.tensor([2, 1])) == torch.tensor([eps + 1, 2 - eps]),
torch.polygamma(1, torch.tensor([1.0, 0.5])),
torch.polygamma(2, torch.tensor([1.0, 0.5])),
torch.polygamma(3, torch.tensor([1.0, 0.5])),
torch.polygamma(4, torch.tensor([1.0, 0.5])),
torch.pow(a, 2),
torch.pow(torch.arange(1.0, 5.0), torch.arange(1.0, 5.0)),
torch.rad2deg(
torch.tensor([[3.142, -3.142], [6.283, -6.283],
[1.570, -1.570]])),
torch.randn(4, dtype=torch.cfloat).real,
torch.reciprocal(a),
torch.remainder(torch.tensor([-3.0, -2.0]), 2),
torch.remainder(torch.tensor([1, 2, 3, 4, 5]), 1.5),
torch.round(a),
torch.rsqrt(a),
torch.sigmoid(a),
torch.sign(torch.tensor([0.7, -1.2, 0.0, 2.3])),
torch.sgn(a),
torch.signbit(torch.tensor([0.7, -1.2, 0.0, 2.3])),
torch.sin(a),
torch.sinc(a),
torch.sinh(a),
torch.sqrt(a),
torch.square(a),
torch.sub(torch.tensor((1, 2)), torch.tensor((0, 1)), alpha=2),
torch.tan(a),
torch.tanh(a),
torch.trunc(a),
torch.xlogy(f, g),
torch.xlogy(f, g),
torch.xlogy(f, 4),
torch.xlogy(2, g),
)
# bitwise_not
torch.bitwise_not(t)
# bitwise_and
torch.bitwise_and(t, torch.tensor([1, 0, 3], dtype=torch.int8))
torch.bitwise_and(torch.tensor([True, True, False]),
torch.tensor([False, True, False]))
# bitwise_or
torch.bitwise_or(t, torch.tensor([1, 0, 3], dtype=torch.int8))
torch.bitwise_or(torch.tensor([True, True, False]),
torch.tensor([False, True, False]))
# bitwise_xor
torch.bitwise_xor(t, torch.tensor([1, 0, 3], dtype=torch.int8))
# ceil
torch.ceil(a)
# clamp/clip
torch.clamp(a, min=-0.5, max=0.5)
torch.clamp(a, min=0.5)
torch.clamp(a, max=0.5)
torch.clip(a, min=-0.5, max=0.5)
# conj
torch.conj(torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]))
# copysign
torch.copysign(a, 1)
def __threefry32(
x0: torch.Tensor, x1: torch.Tensor, seed: int
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Counter-based pseudo random number generator. Based on a 12-round Threefry "encryption" algorithm [1]. Returns
Two vectors with num_samples / 2 (rounded-up) pseudo random numbers. This is the 32-bit version.
Parameters
----------
x0 : torch.Tensor
Upper bits of the to be encoded random sequence
x1 : torch.Tensor
Lower bits of the to be encoded random sequence
seed : int
The seed, i.e. key, for the threefry32 encryption
References
----------
[1] Salmon, John K., Moraes, Mark A., Dror, Ron O. and Shaw, David E., "Parallel random numbers: as easy as 1, 2,
3", Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis,
p. 16, 2011
"""
samples = len(x0)
# Seed is > 32 bit
seed_32 = seed & 0x7FFFFFFF
# set up key buffer
ks_0 = torch.full((samples,), seed_32, dtype=torch.int32, device=x0.device)
ks_1 = torch.full((samples,), seed_32, dtype=torch.int32, device=x1.device)
ks_2 = torch.full((samples,), 466688986, dtype=torch.int32, device=x0.device)
# ks_2 ^= ks_0
# ks_2 ^= ks_1
ks_2 = torch.bitwise_xor(torch.bitwise_xor(ks_2, ks_0), ks_1)
# initialize output using the key
x0 += ks_0
x1 += ks_1
# perform rounds
# round 1
x0 += x1
x1 = (x1 << 13) | ((x1 >> 19) & 0x1FFF)
x1 = torch.bitwise_xor(x1, x0)
# x1 ^= x0
# round 2
x0 += x1
x1 = (x1 << 15) | ((x1 >> 17) & 0x7FFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 3
x0 += x1
x1 = (x1 << 26) | ((x1 >> 6) & 0x3FFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 4
x0 += x1
x1 = (x1 << 6) | ((x1 >> 26) & 0x3F)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# inject key
x0 += ks_1
x1 += ks_2 + 1
# round 5
x0 += x1
x1 = (x1 << 17) | ((x1 >> 15) & 0x1FFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 6
x0 += x1
x1 = (x1 << 29) | ((x1 >> 3) & 0x1FFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 7
x0 += x1
x1 = (x1 << 16) | ((x1 >> 16) & 0xFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 8
x0 += x1
x1 = (x1 << 24) | ((x1 >> 8) & 0xFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# inject key
# x0 += ks_2; x1 += (ks_0 + 2)
#
# x0 += x1; x1 = (x1 << 13) | (x1 >> 19); x1 ^= x0 # round 9
# x0 += x1; x1 = (x1 << 15) | (x1 >> 17); x1 ^= x0 # round 10
# x0 += x1; x1 = (x1 << 26) | (x1 >> 6); x1 ^= x0 # round 11
# x0 += x1; x1 = (x1 << 6) | (x1 >> 26); x1 ^= x0 # round 12
# inject key
x0 += ks_0
x1 += ks_1 + 3
return x0, x1
def b_xor(t, t_):
return Tensor(torch.bitwise_xor(t.tensor, t_.tensor))
def __threefry64(
x0: torch.Tensor, x1: torch.Tensor, seed: int
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Counter-based pseudo random number generator. Based on a 12-round Threefry "encryption" algorithm [1]. This is the
64-bit version.
Parameters
----------
x0 : torch.Tensor
Upper bits of the to be encoded random sequence
x1 : torch.Tensor
Lower bits of the to be encoded random sequence
seed : int
The seed, i.e. key, for the threefry64 encryption
References
----------
[1] Salmon, John K., Moraes, Mark A., Dror, Ron O. and Shaw, David E., "Parallel random numbers: as easy as 1, 2,
3", Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis,
p. 16, 2011
"""
samples = len(x0)
# set up key buffer
ks_0 = torch.full((samples,), seed, dtype=torch.int64, device=x0.device)
ks_1 = torch.full((samples,), seed, dtype=torch.int64, device=x1.device)
ks_2 = torch.full((samples,), 2004413935125273122, dtype=torch.int64, device=x0.device)
# ks_2 ^= ks_0
# ks_2 ^= ks_1
ks_2 = torch.bitwise_xor(torch.bitwise_xor(ks_2, ks_0), ks_1)
# initialize output using the key
x0 += ks_0
x1 += ks_1
# perform rounds
# round 1
x0 += x1
x1 = (x1 << 16) | ((x1 >> 48) & 0xFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 2
x0 += x1
x1 = (x1 << 42) | ((x1 >> 22) & 0x3FFFFFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 3
x0 += x1
x1 = (x1 << 12) | ((x1 >> 52) & 0xFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 4
x0 += x1
x1 = (x1 << 31) | ((x1 >> 33) & 0x7FFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# inject key
x0 += ks_1
x1 += ks_2 + 1
# round 5
x0 += x1
x1 = (x1 << 16) | ((x1 >> 48) & 0xFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 6
x0 += x1
x1 = (x1 << 32) | ((x1 >> 32) & 0xFFFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 7
x0 += x1
x1 = (x1 << 24) | ((x1 >> 40) & 0xFFFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# round 8
x0 += x1
x1 = (x1 << 21) | ((x1 >> 43) & 0x1FFFFF)
# x1 ^= x0
x1 = torch.bitwise_xor(x1, x0)
# inject key
# x0 += ks_2; x1 += (ks_0 + 2)
#
# x0 += x1; x1 = (x1 << 16) | (x1 >> 48); x1 ^= x0 # round 9
# x0 += x1; x1 = (x1 << 42) | (x1 >> 22); x1 ^= x0 # round 10
# x0 += x1; x1 = (x1 << 12) | (x1 >> 52); x1 ^= x0 # round 11
# x0 += x1; x1 = (x1 << 31) | (x1 >> 33); x1 ^= x0 # round 12
# inject key
x0 += ks_0
x1 += ks_1 + 3
return x0, x1
def forward(self, tree_graphs, tree_vec):
device = tree_vec.device
batch_size = tree_graphs.batch_size
root_ids = get_root_ids(tree_graphs)
if 'x' not in tree_graphs.ndata:
tree_graphs.ndata['x'] = self.embedding(tree_graphs.ndata['wid'])
if 'src_x' not in tree_graphs.edata:
tree_graphs.apply_edges(fn.copy_u('x', 'src_x'))
tree_graphs = tree_graphs.local_var()
tree_graphs.apply_edges(func=lambda edges: {'dst_wid': edges.dst['wid']})
line_tree_graphs = dgl.line_graph(tree_graphs, backtracking=False, shared=True)
line_num_nodes = line_tree_graphs.num_nodes()
line_tree_graphs.ndata.update({
'src_x_r': self.W_r(line_tree_graphs.ndata['src_x']),
# Exploit the fact that the reduce function is a sum of incoming messages,
# and uncomputed messages are zero vectors.
'h': torch.zeros(line_num_nodes, self.hidden_size).to(device),
'vec': dgl.broadcast_edges(tree_graphs, tree_vec),
'sum_h': torch.zeros(line_num_nodes, self.hidden_size).to(device),
'sum_gated_h': torch.zeros(line_num_nodes, self.hidden_size).to(device)
})
# input tensors for stop prediction (p) and label prediction (q)
pred_hiddens, pred_mol_vecs, pred_targets = [], [], []
stop_hiddens, stop_targets = [], []
# Predict root
pred_hiddens.append(torch.zeros(batch_size, self.hidden_size).to(device))
pred_targets.append(tree_graphs.ndata['wid'][root_ids.to(device)])
pred_mol_vecs.append(tree_vec)
# Traverse the tree and predict on children
for eid, p in dfs_order(tree_graphs, root_ids.to(dtype=tree_graphs.idtype)):
eid = eid.to(device)
p = p.to(device=device, dtype=tree_graphs.idtype)
# Message passing excluding the target
line_tree_graphs.pull(v=eid, message_func=fn.copy_u('h', 'h_nei'),
reduce_func=fn.sum('h_nei', 'sum_h'))
line_tree_graphs.pull(v=eid, message_func=self.gru_message,
reduce_func=fn.sum('m', 'sum_gated_h'))
line_tree_graphs.apply_nodes(self.gru_update, v=eid)
# Node aggregation including the target
# By construction, the edges of the raw graph follow the order of
# (i1, j1), (j1, i1), (i2, j2), (j2, i2), ... The order of the nodes
# in the line graph corresponds to the order of the edges in the raw graph.
eid = eid.long()
reverse_eid = torch.bitwise_xor(eid, torch.tensor(1).to(device))
cur_o = line_tree_graphs.ndata['sum_h'][eid] + \
line_tree_graphs.ndata['h'][reverse_eid]
# Gather targets
mask = (p == torch.tensor(0).to(device))
pred_list = eid[mask]
stop_target = torch.tensor(1).to(device) - p
# Hidden states for stop prediction
stop_hidden = torch.cat([line_tree_graphs.ndata['src_x'][eid],
cur_o, line_tree_graphs.ndata['vec'][eid]], dim=1)
stop_hiddens.append(stop_hidden)
stop_targets.extend(stop_target)
#Hidden states for clique prediction
if len(pred_list) > 0:
pred_mol_vecs.append(line_tree_graphs.ndata['vec'][pred_list])
pred_hiddens.append(line_tree_graphs.ndata['h'][pred_list])
pred_targets.append(line_tree_graphs.ndata['dst_wid'][pred_list])
#Last stop at root
root_ids = root_ids.to(device)
cur_x = tree_graphs.ndata['x'][root_ids]
tree_graphs.edata['h'] = line_tree_graphs.ndata['h']
tree_graphs.pull(v=root_ids.to(dtype=tree_graphs.idtype),
message_func=fn.copy_e('h', 'm'), reduce_func=fn.sum('m', 'cur_o'))
stop_hidden = torch.cat([cur_x, tree_graphs.ndata['cur_o'][root_ids], tree_vec], dim=1)
stop_hiddens.append(stop_hidden)
stop_targets.extend(torch.zeros(batch_size).to(device))
# Predict next clique
pred_hiddens = torch.cat(pred_hiddens, dim=0)
pred_mol_vecs = torch.cat(pred_mol_vecs, dim=0)
pred_vecs = torch.cat([pred_hiddens, pred_mol_vecs], dim=1)
pred_vecs = F.relu(self.W(pred_vecs))
pred_scores = self.W_o(pred_vecs)
pred_targets = torch.cat(pred_targets, dim=0)
pred_loss = self.pred_loss(pred_scores, pred_targets) / batch_size
_, preds = torch.max(pred_scores, dim=1)
pred_acc = torch.eq(preds, pred_targets).float()
pred_acc = torch.sum(pred_acc) / pred_targets.nelement()
# Predict stop
stop_hiddens = torch.cat(stop_hiddens, dim=0)
stop_vecs = F.relu(self.U(stop_hiddens))
stop_scores = self.U_s(stop_vecs).squeeze()
stop_targets = torch.Tensor(stop_targets).to(device)
stop_loss = self.stop_loss(stop_scores, stop_targets) / batch_size
stops = torch.ge(stop_scores, 0).float()
stop_acc = torch.eq(stops, stop_targets).float()
stop_acc = torch.sum(stop_acc) / stop_targets.nelement()
return pred_loss, stop_loss, pred_acc.item(), stop_acc.item()
