diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index 3abcbb0..c6e57bc 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -1073,6 +1073,34 @@
 #endif  // CONFIG_MOTION_VAR
 #endif  // CONFIG_MOTION_VAR || CONFIG_WARPED_MOTION
 
+// Returns sub-sampled dimensions of the given block.
+// The output values for 'rows_within_bounds' and 'cols_within_bounds' will
+// differ from 'height' and 'width' when part of the block is outside the right
+// and/or bottom image boundary.
+static INLINE void av1_get_block_dimensions(BLOCK_SIZE bsize, int plane,
+                                            const MACROBLOCKD *xd, int *width,
+                                            int *height,
+                                            int *rows_within_bounds,
+                                            int *cols_within_bounds) {
+  const int block_height = block_size_high[bsize];
+  const int block_width = block_size_wide[bsize];
+  const int block_rows = (xd->mb_to_bottom_edge >= 0)
+                             ? block_height
+                             : (xd->mb_to_bottom_edge >> 3) + block_height;
+  const int block_cols = (xd->mb_to_right_edge >= 0)
+                             ? block_width
+                             : (xd->mb_to_right_edge >> 3) + block_width;
+  const struct macroblockd_plane *const pd = &xd->plane[plane];
+  assert(IMPLIES(plane == PLANE_TYPE_Y, pd->subsampling_x == 0));
+  assert(IMPLIES(plane == PLANE_TYPE_Y, pd->subsampling_y == 0));
+  assert(block_width >= block_cols);
+  assert(block_height >= block_rows);
+  if (width) *width = block_width >> pd->subsampling_x;
+  if (height) *height = block_height >> pd->subsampling_y;
+  if (rows_within_bounds) *rows_within_bounds = block_rows >> pd->subsampling_y;
+  if (cols_within_bounds) *cols_within_bounds = block_cols >> pd->subsampling_x;
+}
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif
diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index 8e3debf..d09ce40 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -1266,9 +1266,9 @@
 #endif  // CONFIG_LOOP_RESTORATION
 
 #if CONFIG_PALETTE
-int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
-                                  int c, int n, uint8_t *color_order,
-                                  int *color_idx) {
+int av1_get_palette_color_context(const uint8_t *color_map, int width,
+                                  int stride, int r, int c, int palette_size,
+                                  uint8_t *color_order, int *color_idx) {
   int i;
   // The +10 below should not be needed. But we get a warning "array subscript
   // is above array bounds [-Werror=array-bounds]" without it, possibly due to
@@ -1279,15 +1279,15 @@
   int color_ctx;
   int color_neighbors[4];
   int inverse_color_order[PALETTE_MAX_SIZE];
-  assert(n <= PALETTE_MAX_SIZE);
+  assert(palette_size <= PALETTE_MAX_SIZE);
 
   // Get color indices of neighbors.
-  color_neighbors[0] = (c - 1 >= 0) ? color_map[r * cols + c - 1] : -1;
+  color_neighbors[0] = (c - 1 >= 0) ? color_map[r * stride + c - 1] : -1;
   color_neighbors[1] =
-      (c - 1 >= 0 && r - 1 >= 0) ? color_map[(r - 1) * cols + c - 1] : -1;
-  color_neighbors[2] = (r - 1 >= 0) ? color_map[(r - 1) * cols + c] : -1;
-  color_neighbors[3] = (r - 1 >= 0 && c + 1 <= cols - 1)
-                           ? color_map[(r - 1) * cols + c + 1]
+      (c - 1 >= 0 && r - 1 >= 0) ? color_map[(r - 1) * stride + c - 1] : -1;
+  color_neighbors[2] = (r - 1 >= 0) ? color_map[(r - 1) * stride + c] : -1;
+  color_neighbors[3] = (r - 1 >= 0 && c + 1 <= width - 1)
+                           ? color_map[(r - 1) * stride + c + 1]
                            : -1;
 
   for (i = 0; i < PALETTE_MAX_SIZE; ++i) {
@@ -1306,7 +1306,7 @@
     int max = scores[i];
     int max_idx = i;
     int j;
-    for (j = i + 1; j < n; ++j) {
+    for (j = i + 1; j < palette_size; ++j) {
       if (scores[j] > max) {
         max = scores[j];
         max_idx = j;
@@ -1343,7 +1343,7 @@
   }
 
   if (color_idx != NULL) {
-    *color_idx = inverse_color_order[color_map[r * cols + c]];
+    *color_idx = inverse_color_order[color_map[r * stride + c]];
   }
   return color_ctx;
 }
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index a85202a..c566b73 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -436,9 +436,12 @@
 }
 
 #if CONFIG_PALETTE
-int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
-                                  int c, int n, uint8_t *color_order,
-                                  int *color_idx);
+// Returns the context for palette color index at row 'r' and column 'c',
+// along with the 'color_order' of neighbors and the 'color_idx'.
+// The 'color_map' is a 2D array with the given 'width' and 'stride'.
+int av1_get_palette_color_context(const uint8_t *color_map, int width,
+                                  int stride, int r, int c, int palette_size,
+                                  uint8_t *color_order, int *color_idx);
 #endif  // CONFIG_PALETTE
 
 #ifdef __cplusplus
diff --git a/av1/common/reconintra.c b/av1/common/reconintra.c
index 5545c18..4c260fa 100644
--- a/av1/common/reconintra.c
+++ b/av1/common/reconintra.c
@@ -1925,7 +1925,7 @@
     const int bs = tx_size_wide[tx_size];
     const int stride = wpx;
     int r, c;
-    uint8_t *map = NULL;
+    const uint8_t *const map = xd->plane[plane != 0].color_index_map;
 #if CONFIG_AOM_HIGHBITDEPTH
     uint16_t *palette = xd->mi[0]->mbmi.palette_mode_info.palette_colors +
                         plane * PALETTE_MAX_SIZE;
@@ -1934,8 +1934,6 @@
                        plane * PALETTE_MAX_SIZE;
 #endif  // CONFIG_AOM_HIGHBITDEPTH
 
-    map = xd->plane[plane != 0].color_index_map;
-
 #if CONFIG_AOM_HIGHBITDEPTH
     if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) {
       uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);
diff --git a/av1/decoder/detokenize.c b/av1/decoder/detokenize.c
index eef8a48..50b24e0 100644
--- a/av1/decoder/detokenize.c
+++ b/av1/decoder/detokenize.c
@@ -314,28 +314,35 @@
                                aom_reader *r) {
   const MODE_INFO *const mi = xd->mi[0];
   const MB_MODE_INFO *const mbmi = &mi->mbmi;
-  const BLOCK_SIZE bsize = mbmi->sb_type;
-  const int rows =
-      (block_size_high[bsize]) >> (xd->plane[plane != 0].subsampling_y);
-  const int cols =
-      (block_size_wide[bsize]) >> (xd->plane[plane != 0].subsampling_x);
   uint8_t color_order[PALETTE_MAX_SIZE];
-  const int n = mbmi->palette_mode_info.palette_size[plane != 0];
+  const int n = mbmi->palette_mode_info.palette_size[plane];
   int i, j;
-  uint8_t *color_map = xd->plane[plane != 0].color_index_map;
+  uint8_t *const color_map = xd->plane[plane].color_index_map;
   const aom_prob(*const prob)[PALETTE_COLOR_CONTEXTS][PALETTE_COLORS - 1] =
       plane ? av1_default_palette_uv_color_prob
             : av1_default_palette_y_color_prob;
+  int plane_block_width, plane_block_height, rows, cols;
+  av1_get_block_dimensions(mbmi->sb_type, plane, xd, &plane_block_width,
+                           &plane_block_height, &rows, &cols);
+  assert(plane == 0 || plane == 1);
 
   for (i = 0; i < rows; ++i) {
     for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
-      const int color_ctx = av1_get_palette_color_context(color_map, cols, i, j,
-                                                          n, color_order, NULL);
+      const int color_ctx = av1_get_palette_color_context(
+          color_map, plane_block_width, cols, i, j, n, color_order, NULL);
       const int color_idx = aom_read_tree(r, av1_palette_color_tree[n - 2],
                                           prob[n - 2][color_ctx], ACCT_STR);
       assert(color_idx >= 0 && color_idx < n);
-      color_map[i * cols + j] = color_order[color_idx];
+      color_map[i * plane_block_width + j] = color_order[color_idx];
     }
+    memset(color_map + i * plane_block_width + cols,
+           color_map[i * plane_block_width + cols - 1],
+           (plane_block_width - cols));  // Copy last column to extra columns.
+  }
+  // Copy last row to extra rows.
+  for (i = rows; i < plane_block_height; ++i) {
+    memcpy(color_map + i * plane_block_width,
+           color_map + (rows - 1) * plane_block_width, plane_block_width);
   }
 }
 #endif  // CONFIG_PALETTE
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 02bd950..e1641e3 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1938,14 +1938,14 @@
 #if !CONFIG_PVQ
 #if CONFIG_PALETTE
   for (plane = 0; plane <= 1; ++plane) {
-    if (m->mbmi.palette_mode_info.palette_size[plane] > 0) {
-      const int rows =
-          block_size_high[m->mbmi.sb_type] >> (xd->plane[plane].subsampling_y);
-      const int cols =
-          block_size_wide[m->mbmi.sb_type] >> (xd->plane[plane].subsampling_x);
+    const uint8_t palette_size_plane =
+        m->mbmi.palette_mode_info.palette_size[plane];
+    if (palette_size_plane > 0) {
+      int rows, cols;
+      av1_get_block_dimensions(m->mbmi.sb_type, plane, xd, NULL, NULL, &rows,
+                               &cols);
       assert(*tok < tok_end);
-      pack_palette_tokens(w, tok, m->mbmi.palette_mode_info.palette_size[plane],
-                          rows * cols - 1);
+      pack_palette_tokens(w, tok, palette_size_plane, rows * cols - 1);
       assert(*tok < tok_end + m->mbmi.skip);
     }
   }
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index b9516be..12e9468 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -2270,6 +2270,30 @@
 }
 
 #if CONFIG_PALETTE
+// Extends 'color_map' array from 'orig_width x orig_height' to 'new_width x
+// new_height'. Extra rows and columns are filled in by copying last valid
+// row/column.
+static void extend_palette_color_map(uint8_t *const color_map, int orig_width,
+                                     int orig_height, int new_width,
+                                     int new_height) {
+  int j;
+  assert(new_width >= orig_width);
+  assert(new_height >= orig_height);
+  if (new_width == orig_width && new_height == orig_height) return;
+
+  for (j = orig_height - 1; j >= 0; --j) {
+    memmove(color_map + j * new_width, color_map + j * orig_width, orig_width);
+    // Copy last column to extra columns.
+    memset(color_map + j * new_width + orig_width,
+           color_map[j * new_width + orig_width - 1], new_width - orig_width);
+  }
+  // Copy last row to extra rows.
+  for (j = orig_height; j < new_height; ++j) {
+    memcpy(color_map + j * new_width, color_map + (orig_height - 1) * new_width,
+           new_width);
+  }
+}
+
 static int rd_pick_palette_intra_sby(const AV1_COMP *const cpi, MACROBLOCK *x,
                                      BLOCK_SIZE bsize, int palette_ctx,
                                      int dc_mode_cost, MB_MODE_INFO *best_mbmi,
@@ -2281,12 +2305,13 @@
   MACROBLOCKD *const xd = &x->e_mbd;
   MODE_INFO *const mic = xd->mi[0];
   MB_MODE_INFO *const mbmi = &mic->mbmi;
-  const int rows = block_size_high[bsize];
-  const int cols = block_size_wide[bsize];
   int this_rate, colors, n;
   const int src_stride = x->plane[0].src.stride;
   const uint8_t *const src = x->plane[0].src.buf;
   uint8_t *const color_map = xd->plane[0].color_index_map;
+  int block_width, block_height, rows, cols;
+  av1_get_block_dimensions(bsize, 0, xd, &block_width, &block_height, &rows,
+                           &cols);
 
   assert(cpi->common.allow_screen_content_tools);
 
@@ -2373,6 +2398,8 @@
       pmi->palette_size[0] = k;
 
       av1_calc_indices(data, centroids, color_map, rows * cols, k, 1);
+      extend_palette_color_map(color_map, cols, rows, block_width,
+                               block_height);
       palette_mode_cost =
           dc_mode_cost + cpi->common.bit_depth * k * av1_cost_bit(128, 0) +
           cpi->palette_y_size_cost[bsize - BLOCK_8X8][k - 2] +
@@ -2384,7 +2411,7 @@
         for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
           int color_idx;
           const int color_ctx = av1_get_palette_color_context(
-              color_map, cols, i, j, k, color_order, &color_idx);
+              color_map, cols, block_width, i, j, k, color_order, &color_idx);
           assert(color_idx >= 0 && color_idx < k);
           palette_mode_cost +=
               cpi->palette_y_color_cost[k - 2][color_ctx][color_idx];
@@ -2405,7 +2432,7 @@
       if (this_rd < *best_rd) {
         *best_rd = this_rd;
         memcpy(best_palette_color_map, color_map,
-               rows * cols * sizeof(color_map[0]));
+               block_width * block_height * sizeof(color_map[0]));
         *best_mbmi = *mbmi;
         rate_overhead = this_rate - tokenonly_rd_stats.rate;
         if (rate) *rate = this_rate;
@@ -4339,8 +4366,6 @@
   MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
   PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
   const BLOCK_SIZE bsize = mbmi->sb_type;
-  const int rows = block_size_high[bsize] >> (xd->plane[1].subsampling_y);
-  const int cols = block_size_wide[bsize] >> (xd->plane[1].subsampling_x);
   int this_rate;
   int64_t this_rd;
   int colors_u, colors_v, colors;
@@ -4349,7 +4374,9 @@
   const uint8_t *const src_v = x->plane[2].src.buf;
   uint8_t *const color_map = xd->plane[1].color_index_map;
   RD_STATS tokenonly_rd_stats;
-
+  int plane_block_width, plane_block_height, rows, cols;
+  av1_get_block_dimensions(bsize, 1, xd, &plane_block_width,
+                           &plane_block_height, &rows, &cols);
   if (rows * cols > PALETTE_MAX_BLOCK_SIZE) return;
 
 #if CONFIG_FILTER_INTRA
@@ -4437,6 +4464,8 @@
         centroids[i * 2 + 1] = lb_v + (2 * i + 1) * (ub_v - lb_v) / n / 2;
       }
       av1_k_means(data, centroids, color_map, rows * cols, n, 2, max_itr);
+      extend_palette_color_map(color_map, cols, rows, plane_block_width,
+                               plane_block_height);
       pmi->palette_size[1] = n;
       for (i = 1; i < 3; ++i) {
         for (j = 0; j < n; ++j) {
@@ -4464,8 +4493,9 @@
       for (i = 0; i < rows; ++i) {
         for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
           int color_idx;
-          const int color_ctx = av1_get_palette_color_context(
-              color_map, cols, i, j, n, color_order, &color_idx);
+          const int color_ctx =
+              av1_get_palette_color_context(color_map, cols, plane_block_width,
+                                            i, j, n, color_order, &color_idx);
           assert(color_idx >= 0 && color_idx < n);
           this_rate += cpi->palette_uv_color_cost[n - 2][color_ctx][color_idx];
         }
@@ -4476,7 +4506,8 @@
         *best_rd = this_rd;
         *palette_mode_info = *pmi;
         memcpy(best_palette_color_map, color_map,
-               rows * cols * sizeof(best_palette_color_map[0]));
+               plane_block_width * plane_block_height *
+                   sizeof(best_palette_color_map[0]));
         *mode_selected = DC_PRED;
         *rate = this_rate;
         *distortion = tokenonly_rd_stats.dist;
@@ -9039,8 +9070,6 @@
   MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
   PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
   const BLOCK_SIZE bsize = mbmi->sb_type;
-  const int rows = block_size_high[bsize] >> (xd->plane[1].subsampling_y);
-  const int cols = block_size_wide[bsize] >> (xd->plane[1].subsampling_x);
   int src_stride = x->plane[1].src.stride;
   const uint8_t *const src_u = x->plane[1].src.buf;
   const uint8_t *const src_v = x->plane[2].src.buf;
@@ -9052,6 +9081,9 @@
   const uint16_t *const src_u16 = CONVERT_TO_SHORTPTR(src_u);
   const uint16_t *const src_v16 = CONVERT_TO_SHORTPTR(src_v);
 #endif  // CONFIG_AOM_HIGHBITDEPTH
+  int plane_block_width, plane_block_height, rows, cols;
+  av1_get_block_dimensions(bsize, 1, xd, &plane_block_width,
+                           &plane_block_height, &rows, &cols);
   (void)cpi;
 
   for (r = 0; r < rows; ++r) {
@@ -9078,6 +9110,8 @@
 
   av1_calc_indices(data, centroids, color_map, rows * cols,
                    pmi->palette_size[1], 2);
+  extend_palette_color_map(color_map, cols, rows, plane_block_width,
+                           plane_block_height);
 }
 #endif  // CONFIG_PALETTE
 
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 8cb4bd5..0f3ab7f 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -391,25 +391,26 @@
   const MACROBLOCK *const x = &td->mb;
   const MACROBLOCKD *const xd = &x->e_mbd;
   const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
-  const uint8_t *const color_map = xd->plane[plane != 0].color_index_map;
+  const uint8_t *const color_map = xd->plane[plane].color_index_map;
   const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
-  const int n = pmi->palette_size[plane != 0];
+  const int n = pmi->palette_size[plane];
   int i, j;
   int this_rate = 0;
   uint8_t color_order[PALETTE_MAX_SIZE];
-  const int rows =
-      block_size_high[bsize] >> (xd->plane[plane != 0].subsampling_y);
-  const int cols =
-      block_size_wide[bsize] >> (xd->plane[plane != 0].subsampling_x);
   const aom_prob(*const probs)[PALETTE_COLOR_CONTEXTS][PALETTE_COLORS - 1] =
       plane == 0 ? av1_default_palette_y_color_prob
                  : av1_default_palette_uv_color_prob;
+  int plane_block_width, rows, cols;
+  av1_get_block_dimensions(bsize, plane, xd, &plane_block_width, NULL, &rows,
+                           &cols);
+  assert(plane == 0 || plane == 1);
 
   for (i = 0; i < rows; ++i) {
     for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
       int color_new_idx;
-      const int color_ctx = av1_get_palette_color_context(
-          color_map, cols, i, j, n, color_order, &color_new_idx);
+      const int color_ctx =
+          av1_get_palette_color_context(color_map, cols, plane_block_width, i,
+                                        j, n, color_order, &color_new_idx);
       assert(color_new_idx >= 0 && color_new_idx < n);
       if (dry_run == DRY_RUN_COSTCOEFFS)
         this_rate += cpi->palette_y_color_cost[n - 2][color_ctx][color_new_idx];