code: def _extract_blocks(x, block_h, block_w):
    
    (_, height, width, depth) = common_layers.shape_list(x)
    assert height % block_h == 0
    assert width % block_w == 0
    x = tf.reshape(x, [-1, height//block_h, block_h,
                       width//block_w, block_w, depth])
    return tf.transpose(x, [0, 1, 3, 2, 4, 5])
docs: Helper function for local 2d attention.
  
    Args:
      x: a [batch, height, width, depth] tensor
      block_h: An integer. block height
      block_w: An inteter. block width
  
    returns:
      a [batch, num_heads, height/block_h, width/block_w, depth] tensor