diff --git a/av1/common/txb_common.c b/av1/common/txb_common.c
index b89b40b..a5c8b7d 100644
--- a/av1/common/txb_common.c
+++ b/av1/common/txb_common.c
@@ -282,3 +282,17 @@
 
   return t;
 }
+
+void av1_get_base_level_counts(const uint8_t *const levels,
+                               const int level_minus_1, const int width,
+                               const int height, uint8_t *const level_counts) {
+  const int stride = width + TX_PAD_HOR;
+
+  for (int row = 0; row < height; ++row) {
+    for (int col = 0; col < width; ++col) {
+      level_counts[row * width + col] =
+          get_level_count(levels, stride, row, col, level_minus_1,
+                          base_ref_offset, BASE_CONTEXT_POSITION_NUM);
+    }
+  }
+}
diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 9f040e9..590228b 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -179,14 +179,14 @@
 
 static INLINE int get_base_ctx(const uint8_t *const levels,
                                const int c,  // raster order
-                               const int bwl, const int level_minus_1) {
+                               const int bwl, const int level_minus_1,
+                               const int count) {
   const int row = c >> bwl;
   const int col = c - (row << bwl);
   const int stride = (1 << bwl) + TX_PAD_HOR;
   int mag_count = 0;
   int nb_mag[3] = { 0 };
-  const int count = get_level_count(levels, stride, row, col, level_minus_1,
-                                    base_ref_offset, BASE_CONTEXT_POSITION_NUM);
+
   get_level_mag(levels, stride, row, col, nb_mag);
 
   for (int idx = 0; idx < 3; ++idx)
@@ -567,4 +567,8 @@
 
 void av1_init_lv_map(AV1_COMMON *cm);
 
+void av1_get_base_level_counts(const uint8_t *const levels,
+                               const int level_minus_1, const int width,
+                               const int height, uint8_t *const level_counts);
+
 #endif  // AV1_COMMON_TXB_COMMON_H_
diff --git a/av1/decoder/decodetxb.c b/av1/decoder/decodetxb.c
index f9c8710..1c9f3b1 100644
--- a/av1/decoder/decodetxb.c
+++ b/av1/decoder/decodetxb.c
@@ -205,15 +205,17 @@
 #if USE_CAUSAL_BASE_CTX
   update_eob = *eob - 1;
 #else
+  uint8_t level_counts[MAX_TX_SQUARE];
   int i;
   for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+    av1_get_base_level_counts(levels, i, width, height, level_counts);
     for (c = *eob - 1; c >= 0; --c) {
       uint8_t *const level = &levels[get_paded_idx(scan[c], bwl)];
       int ctx;
 
       if (*level <= i) continue;
 
-      ctx = get_base_ctx(levels, scan[c], bwl, i);
+      ctx = get_base_ctx(levels, scan[c], bwl, i, level_counts[scan[c]]);
 
       if (av1_read_record_bin(
               counts, r, ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx], 2,
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 85d8630..5fd2836 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -371,6 +371,7 @@
   }
   update_eob = eob - 1;
 #else
+  uint8_t level_counts[MAX_TX_SQUARE];
   for (int i = 1; i < eob; ++i) {
     c = eob - 1 - i;
     int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
@@ -383,13 +384,14 @@
   }
 
   for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
+    av1_get_base_level_counts(levels, i, width, height, level_counts);
     for (c = eob - 1; c >= 0; --c) {
       const tran_low_t level = abs(tcoeff[scan[c]]);
       int ctx;
 
       if (level <= i) continue;
 
-      ctx = get_base_ctx(levels, scan[c], bwl, i);
+      ctx = get_base_ctx(levels, scan[c], bwl, i, level_counts[scan[c]]);
 
       if (level == i + 1) {
         aom_write_bin(w, 1, ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx],
@@ -2096,14 +2098,16 @@
 
 #if !USE_CAUSAL_BASE_CTX
   // Reverse process order to handle coefficient level and sign.
+  uint8_t level_counts[MAX_TX_SQUARE];
   for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
+    av1_get_base_level_counts(levels, i, width, height, level_counts);
     for (c = eob - 1; c >= 0; --c) {
       const tran_low_t level = abs(tcoeff[scan[c]]);
       int ctx;
 
       if (level <= i) continue;
 
-      ctx = get_base_ctx(levels, scan[c], bwl, i);
+      ctx = get_base_ctx(levels, scan[c], bwl, i, level_counts[scan[c]]);
 
       if (level == i + 1) {
         ++td->counts->coeff_base[txsize_ctx][plane_type][i][ctx][1];