diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index 0cfbece..b8d8533 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -1201,12 +1201,17 @@
 
 void av1_setup_block_planes(MACROBLOCKD *xd, int ss_x, int ss_y);
 
-static INLINE int tx_size_to_depth(TX_SIZE tx_size) {
-  return (int)(tx_size - TX_SIZE_LUMA_MIN);
+static INLINE int tx_size_cat_to_max_depth(int tx_size_cat) {
+  return AOMMIN(tx_size_cat + 1, MAX_TX_DEPTH);
 }
 
-static INLINE TX_SIZE depth_to_tx_size(int depth) {
-  return (TX_SIZE)(depth + TX_SIZE_LUMA_MIN);
+static INLINE int tx_size_to_depth(TX_SIZE tx_size, int tx_size_cat) {
+  return (int)(tx_size_cat + 1 - (int)tx_size);
+}
+
+static INLINE TX_SIZE depth_to_tx_size(int depth, int tx_size_cat) {
+  assert(tx_size_cat + 1 - depth >= 0 && tx_size_cat + 1 - depth < TX_SIZES);
+  return (TX_SIZE)(tx_size_cat + 1 - depth);
 }
 
 static INLINE TX_SIZE av1_get_uv_tx_size(const MB_MODE_INFO *mbmi,
diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index 92bc02e..6f4d862 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -1757,15 +1757,31 @@
 #endif
 
 static const aom_cdf_prob
-    default_tx_size_cdf[MAX_TX_DEPTH][TX_SIZE_CONTEXTS][CDF_SIZE(MAX_TX_DEPTH +
-                                                                 1)] = {
-      { { AOM_CDF2(12800) }, { AOM_CDF2(8448) } },
-      { { AOM_CDF3(2560, 20496) }, { AOM_CDF3(1920, 14091) } },
-      { { AOM_CDF4(384, 17588, 19782) }, { AOM_CDF4(640, 7166, 8466) } },
+    default_tx_size_cdf[MAX_TX_CATS][TX_SIZE_CONTEXTS][CDF_SIZE(MAX_TX_DEPTH +
+                                                                1)] = {
+#if MAX_TX_DEPTH == 2
+      { { AOM_CDF2(19968) }, { AOM_CDF2(24320) } },
+      { { AOM_CDF3(12272, 30172) }, { AOM_CDF3(18677, 30848) } },
+      { { AOM_CDF3(12986, 15180) }, { AOM_CDF3(24302, 25602) } },
 #if CONFIG_TX64X64
-      { { AOM_CDF5(128, 8288, 21293, 26986) },
-        { AOM_CDF5(128, 4208, 10009, 15965) } },
-#endif
+      { { AOM_CDF3(5782, 11475) }, { AOM_CDF3(16803, 22759) } },
+#endif  // CONFIG_TX64X64
+#elif MAX_TX_DEPTH == 3
+      { { AOM_CDF2(19968) }, { AOM_CDF2(24320) } },
+      { { AOM_CDF3(12272, 30172) }, { AOM_CDF3(18677, 30848) } },
+      { { AOM_CDF4(12986, 15180, 32384) }, { AOM_CDF4(24302, 25602, 32128) } },
+#if CONFIG_TX64X64
+      { { AOM_CDF4(5782, 11475, 24480) }, { AOM_CDF4(16803, 22759, 28560) } },
+#endif  // CONFIG_TX64X64
+#else
+      { { AOM_CDF2(19968) }, { AOM_CDF2(24320) } },
+      { { AOM_CDF3(12272, 30172) }, { AOM_CDF3(18677, 30848) } },
+      { { AOM_CDF4(12986, 15180, 32384) }, { AOM_CDF4(24302, 25602, 32128) } },
+#if CONFIG_TX64X64
+      { { AOM_CDF5(5782, 11475, 24480, 32640) },
+        { AOM_CDF5(16803, 22759, 28560, 32640) } },
+#endif  // CONFIG_TX64X64
+#endif  // MAX_TX_DEPTH == 2
     };
 
 static const aom_cdf_prob
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index 5cd8a84..618b86d 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -326,7 +326,7 @@
                               [CDF_SIZE(2 * MAX_ANGLE_DELTA + 1)];
 #endif  // CONFIG_EXT_INTRA_MOD
 
-  aom_cdf_prob tx_size_cdf[MAX_TX_DEPTH][TX_SIZE_CONTEXTS]
+  aom_cdf_prob tx_size_cdf[MAX_TX_CATS][TX_SIZE_CONTEXTS]
                           [CDF_SIZE(MAX_TX_DEPTH + 1)];
   aom_cdf_prob delta_q_cdf[CDF_SIZE(DELTA_Q_PROBS + 1)];
 #if CONFIG_EXT_DELTA_Q
diff --git a/av1/common/enums.h b/av1/common/enums.h
index e883103..71083e9 100644
--- a/av1/common/enums.h
+++ b/av1/common/enums.h
@@ -246,7 +246,9 @@
    one more than the minimum. */
 #define TX_SIZE_CTX_MIN (TX_SIZE_LUMA_MIN + 1)
 
-#define MAX_TX_DEPTH (TX_SIZES - TX_SIZE_CTX_MIN)
+// Maximum tx_size categories
+#define MAX_TX_CATS (TX_SIZES - TX_SIZE_CTX_MIN)
+#define MAX_TX_DEPTH 2
 
 #define MAX_TX_SIZE_LOG2 (5 + CONFIG_TX64X64)
 #define MAX_TX_SIZE (1 << MAX_TX_SIZE_LOG2)
diff --git a/av1/decoder/decodemv.c b/av1/decoder/decodemv.c
index 3a0486b..4aebae4 100644
--- a/av1/decoder/decodemv.c
+++ b/av1/decoder/decodemv.c
@@ -510,13 +510,15 @@
 
 static TX_SIZE read_selected_tx_size(AV1_COMMON *cm, MACROBLOCKD *xd,
                                      int32_t tx_size_cat, aom_reader *r) {
+  const int max_depths = tx_size_cat_to_max_depth(tx_size_cat);
   const int ctx = get_tx_size_context(xd);
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
   (void)cm;
 
   const int depth = aom_read_symbol(r, ec_ctx->tx_size_cdf[tx_size_cat][ctx],
-                                    tx_size_cat + 2, ACCT_STR);
-  const TX_SIZE tx_size = depth_to_tx_size(depth);
+                                    max_depths + 1, ACCT_STR);
+  assert(depth >= 0 && depth <= max_depths);
+  const TX_SIZE tx_size = depth_to_tx_size(depth, tx_size_cat);
   assert(!is_rect_tx(tx_size));
   return tx_size;
 }
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 9f7e314..23b60c5 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -322,18 +322,22 @@
   (void)cm;
   if (block_signals_txsize(bsize)) {
     const TX_SIZE tx_size = mbmi->tx_size;
-    const int is_inter = is_inter_block(mbmi);
     const int tx_size_ctx = get_tx_size_context(xd);
-    const int32_t tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize]
-                                         : intra_tx_size_cat_lookup[bsize];
+    const int32_t tx_size_cat = intra_tx_size_cat_lookup[bsize];
     const TX_SIZE coded_tx_size = txsize_sqr_up_map[tx_size];
-    const int depth = tx_size_to_depth(coded_tx_size);
+    const int depth = tx_size_to_depth(coded_tx_size, tx_size_cat);
+    const int max_depths = tx_size_cat_to_max_depth(tx_size_cat);
+
+    assert(coded_tx_size <= tx_size_cat + 1);
+    assert(depth >= 0 && depth <= max_depths);
+
+    assert(!is_inter_block(mbmi));
     assert(IMPLIES(is_rect_tx(tx_size), is_rect_tx_allowed(xd, mbmi)));
 
     aom_write_symbol(w, depth, ec_ctx->tx_size_cdf[tx_size_cat][tx_size_ctx],
-                     tx_size_cat + 2);
+                     max_depths + 1);
 #if CONFIG_RECT_TX_EXT
-    if (is_quarter_tx_allowed(xd, mbmi, is_inter) && tx_size != coded_tx_size)
+    if (is_quarter_tx_allowed(xd, mbmi, 0) && tx_size != coded_tx_size)
 #if CONFIG_NEW_MULTISYMBOL
       aom_write_symbol(w, tx_size == quarter_txsize_lookup[bsize],
                        cm->fc->quarter_tx_size_cdf, 2);
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index a44bc4a..6191136 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -219,7 +219,7 @@
   }
 #endif  // CONFIG_CFL
 
-  for (i = 0; i < MAX_TX_DEPTH; ++i)
+  for (i = 0; i < MAX_TX_CATS; ++i)
     for (j = 0; j < TX_SIZE_CONTEXTS; ++j)
       av1_cost_tokens_from_cdf(x->tx_size_cost[i][j], fc->tx_size_cdf[i][j],
                                NULL);
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index d3090c4..08f052d 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -2357,7 +2357,7 @@
     const int32_t tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize]
                                          : intra_tx_size_cat_lookup[bsize];
     const TX_SIZE coded_tx_size = txsize_sqr_up_map[tx_size];
-    const int depth = tx_size_to_depth(coded_tx_size);
+    const int depth = tx_size_to_depth(coded_tx_size, tx_size_cat);
     const int tx_size_ctx = get_tx_size_context(xd);
     int r_tx_size = x->tx_size_cost[tx_size_cat][tx_size_ctx][depth];
 #if CONFIG_RECT_TX_EXT
@@ -2739,7 +2739,7 @@
 
   if (tx_select) {
     start_tx = max_tx_size;
-    end_tx = (max_tx_size >= TX_32X32) ? TX_8X8 : TX_4X4;
+    end_tx = AOMMAX((int)TX_4X4, start_tx - MAX_TX_DEPTH + evaluate_rect_tx);
   } else {
     const TX_SIZE chosen_tx_size =
         tx_size_from_tx_mode(bs, cm->tx_mode, is_inter);