def compute_streaming_mask(chunk_size, history_size, query, value=None):
    """
    Computes a streaming mask as in http://arxiv.org/abs/2010.11395
    For example, if query and value both contain sequences of length 8, chunk size 2, history_size 2
    Chunk size = 2 -> it can see < 2 frames in the future because it in the same chunk, the 2nd frame is the last frame in the chunk therefore it does not see future # pylint: disable=line-too-long
    History size = 2 -> it can see history_size = 2 frames in the past
    The * indicates the current frame
    All frames in the same chunk can see each other
    this function returns a boolean `Tensor` equal to:
    ```
    [[[ 1*,  1,  0,  0,  0,  0,  0,  0 ],
      [ 1,  1*,  0,  0,  0,  0,  0,  0 ],
      [ 1,  1,  1*,  1,  0,  0,  0,  0 ],
      [ 1,  1,  1,  1*,  0,  0,  0,  0 ],
      [ 0,  0,  1,  1,  1*,  1,  0,  0 ],
      [ 0,  0,  1,  1,  1,  1*,  0,  0 ],
      [ 0,  0,  0,  0,  1,  1,  1*,  1 ],
      [ 0,  0,  0,  0,  1,  1,  1,  1*]]]
    ```
    Args:
      chunk_size: chunk size to split
      history_size: history size to keep
      query: query `Tensor` of shape `(B, T, ...)`.
      value: value `Tensor` of shape `(B, S, ...)` (optional, defaults to query).
    Returns:
      mask: a boolean `Tensor` of shape [1, T, S]
    """
    q_seq_length = shape_util.shape_list(query)[1]
    v_seq_length = q_seq_length if value is None else shape_util.shape_list(value)[1]
    hist_size = tf.where(tf.less(history_size, 0), v_seq_length, tf.constant(history_size, tf.int32))

    def _fn(x):
        index = x * chunk_size
        start_index = tf.maximum(0, index - hist_size)
        end_index_excluded = tf.minimum(v_seq_length, index + chunk_size)
        keep = tf.sequence_mask(end_index_excluded, v_seq_length, dtype=tf.bool)
        drop = tf.math.logical_not(tf.sequence_mask(start_index, v_seq_length, dtype=tf.bool))
        return keep & drop

    return tf.expand_dims(tf.map_fn(_fn, tf.math.floordiv(tf.range(q_seq_length, dtype=tf.int32), chunk_size), dtype=tf.bool), axis=0)

    Computes a streaming mask as in http://arxiv.org/abs/2010.11395
    For example, if query and value both contain sequences of length 8, chunk size 2, history_size 2
    Chunk size = 2 -> it can see < 2 frames in the future because it in the same chunk, the 2nd frame is the last frame in the chunk therefore it does not see future # pylint: disable=line-too-long
    History size = 2 -> it can see history_size = 2 frames in the past
    The * indicates the current frame
    All frames in the same chunk can see each other
    this function returns a boolean `Tensor` equal to:
    ```
    [[[ 1*,  1,  0,  0,  0,  0,  0,  0 ],
      [ 1,  1*,  0,  0,  0,  0,  0,  0 ],
      [ 1,  1,  1*,  1,  0,  0,  0,  0 ],
      [ 1,  1,  1,  1*,  0,  0,  0,  0 ],
      [ 0,  0,  1,  1,  1*,  1,  0,  0 ],
      [ 0,  0,  1,  1,  1,  1*,  0,  0 ],
      [ 0,  0,  0,  0,  1,  1,  1*,  1 ],
      [ 0,  0,  0,  0,  1,  1,  1,  1*]]]
    ```
    Args:
      chunk_size: chunk size to split
      history_size: history size to keep
      query: query `Tensor` of shape `(B, T, ...)`.
      value: value `Tensor` of shape `(B, S, ...)` (optional, defaults to query).
    Returns:
      mask: a boolean `Tensor` of shape [1, T, S]
    
tensorflow_asr/models/layers/multihead_attention.py