code: def reshape_by_blocks(x, x_shape, memory_block_size):
    
    x = tf.reshape(x, [
        x_shape[0], x_shape[1], x_shape[2] // memory_block_size,
        memory_block_size, x_shape[3]
    ])
    return x
docs: Reshapes input by splitting its length over blocks of memory_block_size.
  
    Args:
      x: a Tensor with shape [batch, heads, length, depth]
      x_shape: tf.TensorShape of x.
      memory_block_size: Integer which divides length.
  
    Returns:
      Tensor with shape
      [batch, heads, length // memory_block_size, memory_block_size, depth].