in blocksparse/transformer.py [0:0]
def masked_softmax_test(self, x, scale=1.0, autoregress_at_key=None):
y = np.empty_like(x)
m = self.softmax_mask_np # heads, blocks, blk_size
bsize = self.blk_size
ones = (1 << bsize) - 1
for n in range(x.shape[0]):
for h in range(x.shape[1]):
hl = h if self.lut_heads > 1 else 0
for lut in self.nn_list[hl]:
xm = np.full((len(lut), bsize * bsize), -np.finfo(np.float32).max, dtype=np.float32)
for i, (b, k) in enumerate(lut):
xb = x[n,h,b,:,:].reshape(-1)
if m is None:
# apply scale
xm[i,:] = xb * scale
else:
mask = m[hl,b,:]
if autoregress_at_key is not None:
Q = self.nt_list[hl][b][0] * bsize
K = k * bsize
new_mask = np.empty(bsize, dtype=mask.dtype)
for q in range(bsize):
shift_a = bsize - min(max(autoregress_at_key - K, 0), bsize)
shift_b = min(max(bsize-1 + K - (Q + q), 0), bsize)
shift_c = int(min(shift_a, shift_b))
#print(ones, shift_c, type(shift_c))
new_mask[q] = int(mask[q]) & (ones >> shift_c)
mask = new_mask
# apply mask and scale to x block
mask = np.unpackbits(mask.view(np.uint8)).reshape(-1,8)[:,::-1].reshape(-1)
nzIdx = np.nonzero(mask)
xm[i,nzIdx] = xb[nzIdx] * scale
# compute softmax for collection of k blocks
xm = xm.reshape((len(lut), bsize, bsize))
xm = np.exp(xm - np.max(xm, axis=(0,2), keepdims=True))
ym = xm / np.sum(xm, axis=(0,2), keepdims=True)
for i, (b, k) in enumerate(lut):
y[n,h,b,:,:] = ym[i]
return y