def construct(self, inputs, mask=None):
r"""Compute layer output.
Args:
input (torch.Tensor): input data.
mask (torch.Tensor, optional): mask to be applied; e.g. neighbors mask.
Returns:
torch.Tensor: layer output.
"""
# mask input
if mask is not None:
inputs = inputs * F.expand_dims(mask, -1)
# compute sum of input along axis
y = self.reduce_sum(inputs, self.axis)
# compute average of input along axis
if self.average:
# get the number of items along axis
if mask is not None:
N = self.reduce_sum(mask, self.axis)
N = self.maximum(N, other=F.ones_like(N))
else:
N = inputs.shape[self.axis]
y = y / N
return y
def construct(self, atom_types):
# [B,1,1]
exones = self.ones((atom_types.shape[0], 1, 1), ms.int32)
# broadcast to [B*A*N]: [B,1,1] * [1,A,N]
exnfc = exones * F.expand_dims(self.nfc, 0)
exnnc = exones * F.expand_dims(self.nnc, 0)
tmask = F.select(atom_types > 0, F.ones_like(atom_types),
F.ones_like(atom_types) * -1)
tmask = F.cast(tmask, ms.float32)
extmask = F.expand_dims(tmask, -1) * self.nones
mask0 = F.gather(tmask, self.ar0, -1)
mask0 = F.expand_dims(mask0, -2) * self.eaones
mask1 = F.gather(tmask, self.ar1, -1)
mask1 = F.expand_dims(mask1, -2) * self.eaones
mtmp = F.select(exnfc > exnnc, mask1, mask0)
mask = F.select(extmask > 0, mtmp, F.ones_like(mtmp) * -1)
mask = mask > 0
idx = F.select(mask, exnfc, exnnc)
return idx, mask
def construct(self,
positions,
neighbors,
neighbor_mask=None,
cell=None,
cell_offsets=None):
r"""Compute distance of every atom to its neighbors.
Args:
positions (ms.Tensor[float]): atomic Cartesian coordinates with
(N_b x N_at x 3) shape.
neighbors (ms.Tensor[int]): indices of neighboring atoms to consider
with (N_b x N_at x N_nbh) or (N_at x N_nbh) shape.
cell (ms.tensor[float], optional): periodic cell of (N_b x 3 x 3) shape.
cell_offsets (ms.Tensor[float], optional): offset of atom in cell coordinates
with (N_b x N_at x N_nbh x 3) shape.
neighbor_mask (ms.Tensor[bool], optional): boolean mask for neighbor
positions. Required for the stable computation of forces in
molecules with different sizes.
Returns:
ms.Tensor[float]: layer output of (N_b x N_at x N_nbh) shape.
"""
pos_xyz = self.gather_neighbors(positions, neighbors)
# Subtract positions of central atoms to get distance vectors
dist_vec = pos_xyz - F.expand_dims(positions, -2)
# distances = self.norm(dist_vec)
distances = F.square(dist_vec)
distances = self.reducesum(distances, -1)
distances = self.pow(distances, 0.5)
if neighbor_mask is not None:
distances = F.select(neighbor_mask, distances,
F.ones_like(distances) * 999)
return distances
def __init__(self,
d_min=0.0,
d_max=5.0,
num_rbf=32,
sigma=None,
centered=False,
trainable=False):
super().__init__()
# compute offset and width of Gaussian functions
offset = Tensor(np.linspace(d_min, d_max, num_rbf), ms.float32)
if sigma is None:
sigma = (d_max - d_min) / (num_rbf - 1)
width = sigma * F.ones_like(offset)
self.width = width
self.offset = offset
self.centered = centered
if trainable:
self.width = ms.Parameter(width, "widths")
self.offset = ms.Parameter(offset, "offset")
def construct(self, num):
nmax = num * self.ones
idx = F.ones_like(num) * self.range()
return idx < nmax
def construct(self, query, key, value, cutoff=None, mask=None):
r"""Compute multi-head attention.
Args:
query (Mindspore.Tensor [B, A, 1, V]):
key (Mindspore.Tensor [B, A, N', V]):
value (Mindspore.Tensor [B, A, N', V]):
cutoff (Mindspore.Tensor [B, A, 1, N'] or [B, A, 1, 1, N']):
Returns:
Mindspore.Tensor [B, A, V]: multi-head attention output.
"""
if self.n_heads > 1:
q_reshape = query.shape[:-1] + self.reshape_tail
k_reshape = key.shape[:-1] + self.reshape_tail
v_reshape = value.shape[:-1] + self.reshape_tail
# [B, A, 1, h, v]
Q = F.reshape(query, q_reshape)
# [B, A, h, 1, v]
Q = self.transpose(Q, self.trans_shape)
# [B, A, N', h, v]
K = F.reshape(key, k_reshape)
# [B, A, h, N', v]
K = self.transpose(K, self.trans_shape)
# [B, A, N', h, v]
V = F.reshape(value, v_reshape)
# [B, A, h, N', v]
V = self.transpose(V, self.trans_shape)
# [B, A, h, 1, v] x [B, A, h, N', v]^T / \sqrt(v)
# [B, A, h, 1, v] x [B, A, h, v, N'] = [B, A, h, 1, N']
attention_scores = self.bmmt(Q, K)
attention_scores = self.mul(attention_scores, self.scores_mul)
if cutoff is None:
attention_probs = self.softmax(attention_scores)
else:
# [B, A, 1, 1, N']
exmask = F.expand_dims(F.expand_dims(mask, -2), -2)
# [B, A, h, 1, N']
mhmask = exmask * self.exones
large_neg = F.ones_like(attention_scores) * -5e4
attention_scores = F.select(mhmask > 0, attention_scores,
large_neg)
attention_probs = self.softmax(attention_scores)
excut = F.expand_dims(F.expand_dims(cutoff, -2), -2)
# [B, A, h, 1, N'] * [B, A, 1, 1, N']
attention_probs = self.mul(attention_probs, excut)
# [B, A, h, 1, N'] x [B, A, h, N', v] = [B, A, h, 1, v]
context = self.bmm(attention_probs, V)
# [B, A, 1, h, v]
context = self.transpose(context, self.trans_shape)
# [B, A, 1, V]
context = F.reshape(context, query.shape)
else:
# [B, A, 1, V] x [B, A, N', V]^T / \sqrt(V)
# [B, A, 1, V] x [B, A, V, N'] = [B, A, 1, N']
attention_scores = self.bmmt(query, key) * self.scores_mul
if cutoff is None:
attention_probs = self.softmax(attention_scores)
else:
large_neg = F.ones_like(attention_scores) * -5e4
attention_scores = F.select(mask, attention_scores, large_neg)
attention_probs = self.softmax(attention_scores)
# [B, A, 1, N'] * [B, A, 1, N']
attention_probs = attention_probs * F.expand_dims(cutoff, -2)
# [B, A, 1, N'] x [B, A, N', V] = [B, A, 1, V]
context = self.bmm(attention_probs, value)
# [B, A, V]
context = self.squeeze(context)
return self.output(context)
def __init__(self, mean, stddev, eps=1e-9):
super().__init__()
self.mean = mean
self.stddev = stddev
self.eps = F.ones_like(stddev) * eps