Skip to content

Commit

Permalink
initial streaming implementation for chunkwise attention. This is the…
Browse files Browse the repository at this point in the history
… most basic step, where no context is used, therefore it becomes a stateless attention.

PiperOrigin-RevId: 454905548
  • Loading branch information
yqwangustc authored and copybara-github committed Jun 14, 2022
1 parent c86ea38 commit 4121222
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions lingvo/core/batch_major_attention.py
Original file line number Diff line number Diff line change
Expand Up @@ -3110,6 +3110,16 @@ def ExtendStep(self,
msg = 'Not implemented yet'
raise NotImplementedError(msg)

def zero_state(self, batch_size):
"""Returns the initial state given the batch size."""
del batch_size
p = self.params
if p.left_context != 1 or p.right_context != 0:
msg = ('Streaming implementation of chunkwise attention with left context'
'or right context is not supported yet')
raise NotImplementedError(msg)
return py_utils.NestedMap()


class ChunkwiseSelfAttentionXL(ChunkwiseSelfAttention):
"""Chunkwise Self Attention with relative position embedding."""
Expand Down

0 comments on commit 4121222

Please sign in to comment.