diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index b6e4546faf7a46047a55693740431592886a5f13..ec3dc28cfc665722342786392262a36ef77b9b24 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -2528,13 +2528,14 @@ void av1_adapt_intra_frame_probs(AV1_COMMON *cm) {
                          counts->uv_mode[i], fc->uv_mode_prob[i]);
 
 #if CONFIG_EXT_PARTITION_TYPES
-  aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[0],
-                       counts->partition[0], fc->partition_prob[0]);
-  for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; i++)
+  for (i = 0; i < PARTITION_PLOFFSET; ++i)
+    aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[i],
+                         counts->partition[i], fc->partition_prob[i]);
+  for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     aom_tree_merge_probs(av1_ext_partition_tree, pre_fc->partition_prob[i],
                          counts->partition[i], fc->partition_prob[i]);
 #else
-  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; i++) {
+  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i) {
     aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[i],
                          counts->partition[i], fc->partition_prob[i]);
   }
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index b6bc8742185e4cf773cf7be380442c706ecb2b9e..96afcb514a084f87773382be38c34615c996c4d5 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -4611,9 +4611,10 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
   }
 
 #if CONFIG_EXT_PARTITION_TYPES
-  for (i = 0; i < PARTITION_TYPES - 1; ++i)
-    av1_diff_update_prob(&r, &fc->partition_prob[0][i], ACCT_STR);
-  for (j = 1; j < PARTITION_CONTEXTS_PRIMARY; ++j)
+  for (j = 0; j < PARTITION_PLOFFSET; ++j)
+    for (i = 0; i < PARTITION_TYPES - 1; ++i)
+      av1_diff_update_prob(&r, &fc->partition_prob[j][i], ACCT_STR);
+  for (; j < PARTITION_CONTEXTS_PRIMARY; ++j)
     for (i = 0; i < EXT_PARTITION_TYPES - 1; ++i)
       av1_diff_update_prob(&r, &fc->partition_prob[j][i], ACCT_STR);
 #else
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 12e8a818d31a821e08a2d0264a6c1933e7e1b163..b7b82083e73d1cbb1b27504e530c49d3077e4390 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -4666,17 +4666,17 @@ static uint32_t write_compressed_header(AV1_COMP *cpi, uint8_t *data) {
   }
 
 #if CONFIG_EXT_PARTITION_TYPES
-  prob_diff_update(av1_partition_tree, fc->partition_prob[0],
-                   counts->partition[0], PARTITION_TYPES, probwt, header_bc);
-  for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+  for (i = 0; i < PARTITION_PLOFFSET; ++i)
+    prob_diff_update(av1_partition_tree, fc->partition_prob[i],
+                     counts->partition[i], PARTITION_TYPES, probwt, header_bc);
+  for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_ext_partition_tree, fc->partition_prob[i],
                      counts->partition[i], EXT_PARTITION_TYPES, probwt,
                      header_bc);
 #else
-  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i) {
+  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_partition_tree, fc->partition_prob[i],
                      counts->partition[i], PARTITION_TYPES, probwt, header_bc);
-  }
 #endif  // CONFIG_EXT_PARTITION_TYPES
 #if CONFIG_UNPOISON_PARTITION_CTX
   for (; i < PARTITION_CONTEXTS_PRIMARY + PARTITION_BLOCK_SIZES; ++i) {
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index 76d18ea7c16aef2700f1b53cd9c54c45456f31fd..d6561fc114e8c7c628f6e5bf0d5e7efa606cbdef 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -405,9 +405,10 @@ void av1_initialize_rd_consts(AV1_COMP *cpi) {
       cpi->partition_cost[0][PARTITION_SPLIT] = 0;
 #endif
 #if CONFIG_EXT_PARTITION_TYPES
-      av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX],
-                      cm->fc->partition_prob[0], av1_partition_tree);
-      for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+      for (i = 0; i < PARTITION_PLOFFSET; ++i)
+        av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
+                        cm->fc->partition_prob[i], av1_partition_tree);
+      for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
         av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
                         cm->fc->partition_prob[i], av1_ext_partition_tree);
 #else