def get_write_head_offset(self, write_head_index):
    """Lookup the offset to shift the write head at each step.

    By default, we move each write head forward by 1.

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

    Args:
      write_head_index: Identifies which write head we're getting the index for.

    Returns:
      An integer offset to move the write head at each step.
    """
    if write_head_index == 0:
      return 1
    else:
      raise ValueError("Write head index must be 0 for stack.")
Lookup the offset to shift the write head at each step.

    By default, we move each write head forward by 1.

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

    Args:
      write_head_index: Identifies which write head we're getting the index for.

    Returns:
      An integer offset to move the write head at each step.
    
tensor2tensor/models/research/neural_stack.py