def test_incremental_state_multihead_attention(self):
module1 = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
module1 = torch.jit.script(module1)
module2 = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
module2 = torch.jit.script(module2)
state = {}
state = module1.set_incremental_state(state, "key", {"a": torch.tensor([1])})
state = module2.set_incremental_state(state, "key", {"a": torch.tensor([2])})
v1 = module1.get_incremental_state(state, "key")["a"]
v2 = module2.get_incremental_state(state, "key")["a"]
self.assertEqual(v1, 1)
self.assertEqual(v2, 2)
def __init__(self,
decoder_hidden_state_dim,
context_dim,
*,
nheads=1,
unseen_mask=False,
src_length_mask=True):
super().__init__(decoder_hidden_state_dim, context_dim)
assert decoder_hidden_state_dim == context_dim
d_model = decoder_hidden_state_dim # for brevity
assert d_model % nheads == 0
if unseen_mask:
raise NotImplementedError(
"Unseen mask not supported with sequential decoding")
self._fair_attn = fair_multihead.MultiheadAttention(d_model, nheads)
self.use_src_length_mask = src_length_mask
def test_export_multihead_attention(self):
module = multihead_attention.MultiheadAttention(embed_dim=8,
num_heads=2)
torch.jit.script(module)
def test_export_multihead_attention(self):
module = multihead_attention.MultiheadAttention(embed_dim=8,
num_heads=2)
scripted = torch.jit.script(module)
_test_save_and_load(scripted)