code: def get_shifted_center_blocks(x, indices):
    
    center_x = gather_blocks_2d(x, indices)
  
    # Shift right along the length dimension
    def shift_right_2d_blocks(x):
      """Shift the second to last dimension of x right by one."""
      shifted_targets = (
          tf.pad(x, [[0, 0], [0, 0], [0, 0], [1, 0], [0, 0]])[:, :, :, :-1, :])
      return shifted_targets
  
    x_shifted = shift_right_2d_blocks(center_x)
    return x_shifted
docs: Get right shifted blocks for masked local attention 2d.
  
    Args:
      x: A tensor with shape [batch, heads, height, width, depth]
      indices: The indices to gather blocks
  
    Returns:
      x_shifted: a tensor of extracted blocks, each block right shifted along
        length.