def get_read_mask(self, read_head_index):
    """Creates a mask which allows us to attenuate subsequent read strengths.

    This is exposed as its own function so that it can be overridden to provide
    alternate read adressing schemes.

    Args:
      read_head_index: Identifies which read head we're getting the mask for.

    Returns:
      A tf.float32 tensor of shape [1, 1, memory_size, memory_size]
    """
    if read_head_index == 0:
      return tf.expand_dims(
          common_layers.mask_pos_lt(self._memory_size, self._memory_size),
          axis=0)
    else:
      raise ValueError("Read head index must be 0 for stack.")
Creates a mask which allows us to attenuate subsequent read strengths.

    This is exposed as its own function so that it can be overridden to provide
    alternate read adressing schemes.

    Args:
      read_head_index: Identifies which read head we're getting the mask for.

    Returns:
      A tf.float32 tensor of shape [1, 1, memory_size, memory_size]
    
tensor2tensor/models/research/neural_stack.py