def test_masked_select_broadcast(test_case):
x = flow.ones(2, 3, 3)
mask = flow.triu(flow.ones(3, 3), 1)
flow_res = flow.masked_select(x, mask)
np_res = [1, 1, 1, 1, 1, 1]
test_case.assertTrue(
np.allclose(flow_res.numpy(), np_res, 1e-05, 1e-05))
def get_subsequent_mask(seq):
""" For masking out the subsequent info. """
sz_b, len_s = seq.size()
subsequent_mask = flow.triu(flow.ones((len_s, len_s),
device=seq.device,
dtype=flow.int8),
diagonal=1)
subsequent_mask = subsequent_mask.unsqueeze(0).expand(sz_b, -1, -1)
return subsequent_mask
def _test_triu(test_case, diagonal, device):
arr_shape = (4, 4, 8)
np_arr = np.random.randn(*arr_shape)
input_tensor = flow.tensor(np_arr,
dtype=flow.float32,
device=flow.device(device),
requires_grad=True)
output = flow.triu(input_tensor, diagonal=diagonal)
np_out = np.triu(np_arr, diagonal)
test_case.assertTrue(np.allclose(output.numpy(), np_out, 1e-06, 1e-06))
output = output.sum()
output.backward()
np_grad = np.triu(np.ones(shape=arr_shape, dtype=np.float32), diagonal)
test_case.assertTrue(
np.allclose(input_tensor.grad.numpy(), np_grad, 1e-06, 1e-06))
def _triu(self, diagonal=0):
return flow.triu(self, diagonal=diagonal)
def generate_subsequent_mask(self, tgt_len, src_len):
mask = flow.triu(flow.ones((tgt_len, src_len)), 1).to(flow.int32)
return mask
def generate_square_subsequent_mask(self, sz: int) -> Tensor:
return flow.triu(flow.ones((sz, sz)), diagonal=1).to(flow.int32)