Make txfm_partition_update support rectangular tx_size Change-Id: I7d2414a8766141d5109b599271179bc505c772d3