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)