diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index d1d6ecc..b844594 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -5007,9 +5007,10 @@
     txfm_partition_update(xd->above_txfm_context + tx_col,
                           xd->left_txfm_context + tx_row, tx_size);
   } else {
-    BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
-    int bh = num_4x4_blocks_high_lookup[bsize];
+    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
+    const int bs = tx_size_wide_unit[sub_txs];
     int i;
+
     ++counts->txfm_partition[ctx][1];
     ++x->txb_split_count;
 
@@ -5022,10 +5023,10 @@
     }
 
     for (i = 0; i < 4; ++i) {
-      int offsetr = (i >> 1) * bh / 2;
-      int offsetc = (i & 0x01) * bh / 2;
-      update_txfm_count(x, xd, counts, tx_size - 1, depth + 1,
-                        blk_row + offsetr, blk_col + offsetc);
+      int offsetr = (i >> 1) * bs;
+      int offsetc = (i & 0x01) * bs;
+      update_txfm_count(x, xd, counts, sub_txs, depth + 1, blk_row + offsetr,
+                        blk_col + offsetc);
     }
   }
 }
@@ -5037,8 +5038,8 @@
   const int mi_width = num_4x4_blocks_wide_lookup[plane_bsize];
   const int mi_height = num_4x4_blocks_high_lookup[plane_bsize];
   TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
-  BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
-  int bh = num_4x4_blocks_wide_lookup[txb_size];
+  const int bh = tx_size_high_unit[max_tx_size];
+  const int bw = tx_size_wide_unit[max_tx_size];
   int idx, idy;
 
   xd->above_txfm_context = cm->above_txfm_context + mi_col;
@@ -5046,7 +5047,7 @@
       xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
 
   for (idy = 0; idy < mi_height; idy += bh)
-    for (idx = 0; idx < mi_width; idx += bh)
+    for (idx = 0; idx < mi_width; idx += bw)
       update_txfm_count(x, xd, td_counts, max_tx_size, mi_width != mi_height,
                         idy, idx);
 }
@@ -5068,8 +5069,8 @@
                           xd->left_txfm_context + tx_row, tx_size);
 
   } else {
-    BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
-    int bsl = b_width_log2_lookup[bsize];
+    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
+    const int bsl = tx_size_wide_unit[sub_txs];
     int i;
 
     if (tx_size == TX_8X8) {
@@ -5081,11 +5082,10 @@
     }
 
     assert(bsl > 0);
-    --bsl;
     for (i = 0; i < 4; ++i) {
-      int offsetr = (i >> 1) << bsl;
-      int offsetc = (i & 0x01) << bsl;
-      set_txfm_context(xd, tx_size - 1, blk_row + offsetr, blk_col + offsetc);
+      int offsetr = (i >> 1) * bsl;
+      int offsetc = (i & 0x01) * bsl;
+      set_txfm_context(xd, sub_txs, blk_row + offsetr, blk_col + offsetc);
     }
   }
 }
@@ -5096,8 +5096,8 @@
   const int mi_width = num_4x4_blocks_wide_lookup[plane_bsize];
   const int mi_height = num_4x4_blocks_high_lookup[plane_bsize];
   TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
-  BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
-  int bh = num_4x4_blocks_wide_lookup[txb_size];
+  const int bh = tx_size_high_unit[max_tx_size];
+  const int bw = tx_size_wide_unit[max_tx_size];
   int idx, idy;
 
   xd->above_txfm_context = cm->above_txfm_context + mi_col;
@@ -5105,7 +5105,7 @@
       xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
 
   for (idy = 0; idy < mi_height; idy += bh)
-    for (idx = 0; idx < mi_width; idx += bh)
+    for (idx = 0; idx < mi_width; idx += bw)
       set_txfm_context(xd, max_tx_size, idy, idx);
 }
 #endif
diff --git a/av1/encoder/encodemb.c b/av1/encoder/encodemb.c
index 3a33f35..8b62a22 100644
--- a/av1/encoder/encodemb.c
+++ b/av1/encoder/encodemb.c
@@ -875,7 +875,9 @@
   if (tx_size == plane_tx_size) {
     encode_block(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
   } else {
-    int bsl = block_size_wide[bsize] >> (tx_size_wide_log2[0] + 1);
+    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
+    // This is the square transform block partition entry point.
+    int bsl = tx_size_wide_unit[sub_txs];
     int i;
     assert(bsl > 0);
 #if CONFIG_EXT_TX
@@ -885,7 +887,6 @@
     for (i = 0; i < 4; ++i) {
       const int offsetr = blk_row + ((i >> 1) * bsl);
       const int offsetc = blk_col + ((i & 0x01) * bsl);
-      const TX_SIZE sub_txs = tx_size - 1;
       int step = tx_size_wide_unit[sub_txs] * tx_size_high_unit[sub_txs];
 
       if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index d79763e..82cf39a 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -572,22 +572,12 @@
   const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
   const int tx_row = blk_row >> (1 - pd->subsampling_y);
   const int tx_col = blk_col >> (1 - pd->subsampling_x);
+  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
+  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
   TX_SIZE plane_tx_size;
 
-  int max_blocks_high = block_size_high[plane_bsize];
-  int max_blocks_wide = block_size_wide[plane_bsize];
-
   assert(tx_size < TX_SIZES);
 
-  if (xd->mb_to_bottom_edge < 0)
-    max_blocks_high += xd->mb_to_bottom_edge >> (3 + pd->subsampling_y);
-  if (xd->mb_to_right_edge < 0)
-    max_blocks_wide += xd->mb_to_right_edge >> (3 + pd->subsampling_x);
-
-  // Scale to the transform block unit.
-  max_blocks_high >>= tx_size_wide_log2[0];
-  max_blocks_wide >>= tx_size_wide_log2[0];
-
   if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
 
   plane_tx_size =
@@ -605,7 +595,8 @@
       cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
   } else {
     // Half the block size in transform block unit.
-    int bsl = block_size_wide[bsize] >> (tx_size_wide_log2[0] + 1);
+    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
+    const int bsl = tx_size_wide_unit[sub_txs];
     int i;
 
     assert(bsl > 0);
@@ -614,8 +605,6 @@
       const int offsetr = blk_row + ((i >> 1) * bsl);
       const int offsetc = blk_col + ((i & 0x01) * bsl);
 
-      // TODO(jingning): Fix this tx_size transition.
-      const TX_SIZE sub_txs = tx_size - 1;
       int step = tx_size_wide_unit[sub_txs] * tx_size_high_unit[sub_txs];
 
       if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;