diff --git a/src/context/block_unit.rs b/src/context/block_unit.rs
index c39554c8fc..9a26c8537f 100644
--- a/src/context/block_unit.rs
+++ b/src/context/block_unit.rs
@@ -7,6 +7,8 @@
 // Media Patent License 1.0 was not distributed with this source code in the
 // PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 
+use std::mem::MaybeUninit;
+
 use super::*;
 
 use crate::predict::PredictionMode;
@@ -1918,10 +1920,11 @@ impl<'a> ContextWriter<'a> {
     w: &mut W,
   ) {
     // SAFETY: We write to the array below before reading from it.
-    let mut coeff_contexts: Aligned<[i8; MAX_CODED_TX_SQUARE]> =
+    let mut coeff_contexts: Aligned<[MaybeUninit<i8>; MAX_CODED_TX_SQUARE]> =
       unsafe { Aligned::uninitialized() };
 
-    self.get_nz_map_contexts(
+    // get_nz_map_contexts sets coeff_contexts contiguously as a parallel array for scan, not in scan order
+    let coeff_contexts = self.get_nz_map_contexts(
       levels,
       scan,
       eob,
@@ -1932,9 +1935,13 @@ impl<'a> ContextWriter<'a> {
 
     let bhl = Self::get_txb_bhl(tx_size);
 
-    for (c, (&pos, &v)) in scan.iter().zip(coeffs.iter()).enumerate().rev() {
+    let scan_with_ctx =
+      scan.iter().copied().zip(coeff_contexts.iter().copied());
+    for (c, ((pos, coeff_ctx), v)) in
+      scan_with_ctx.zip(coeffs.iter().copied()).enumerate().rev()
+    {
       let pos = pos as usize;
-      let coeff_ctx = coeff_contexts.data[pos];
+      let coeff_ctx = coeff_ctx as usize;
       let level = v.abs();
 
       if c == usize::from(eob) - 1 {
@@ -1942,14 +1949,14 @@ impl<'a> ContextWriter<'a> {
           self,
           w,
           cmp::min(u32::cast_from(level), 3) - 1,
-          &self.fc.coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx as usize]
+          &self.fc.coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx]
         );
       } else {
         symbol_with_update!(
           self,
           w,
           cmp::min(u32::cast_from(level), 3),
-          &self.fc.coeff_base_cdf[txs_ctx][plane_type][coeff_ctx as usize]
+          &self.fc.coeff_base_cdf[txs_ctx][plane_type][coeff_ctx]
         );
       }
 
diff --git a/src/context/transform_unit.rs b/src/context/transform_unit.rs
index a7dbd6c50b..a350cb4a3f 100644
--- a/src/context/transform_unit.rs
+++ b/src/context/transform_unit.rs
@@ -11,6 +11,7 @@ use super::*;
 use crate::predict::PredictionMode;
 use crate::predict::PredictionMode::*;
 use crate::transform::TxType::*;
+use std::mem::MaybeUninit;
 
 pub const MAX_TX_SIZE: usize = 64;
 
@@ -905,25 +906,33 @@ impl<'a> ContextWriter<'a> {
     Self::get_nz_map_ctx_from_stats(stats, coeff_idx, bhl, tx_size, tx_class)
   }
 
-  pub fn get_nz_map_contexts(
+  /// `coeff_contexts_no_scan` is not in the scan order.
+  /// Value for `pos = scan[i]` is at `coeff[i]`, not at `coeff[pos]`.
+  pub fn get_nz_map_contexts<'c>(
     &self, levels: &mut [u8], scan: &[u16], eob: u16, tx_size: TxSize,
-    tx_class: TxClass, coeff_contexts: &mut [i8],
-  ) {
+    tx_class: TxClass, coeff_contexts_no_scan: &'c mut [MaybeUninit<i8>],
+  ) -> &'c mut [i8] {
     let bhl = Self::get_txb_bhl(tx_size);
     let area = av1_get_coded_tx_size(tx_size).area();
-    for i in 0..eob {
-      let pos = scan[i as usize];
-      coeff_contexts[pos as usize] = Self::get_nz_map_ctx(
+
+    let scan = &scan[..usize::from(eob)];
+    let coeffs = &mut coeff_contexts_no_scan[..usize::from(eob)];
+    for (i, (coeff, pos)) in
+      coeffs.iter_mut().zip(scan.iter().copied()).enumerate()
+    {
+      coeff.write(Self::get_nz_map_ctx(
         levels,
         pos as usize,
         bhl,
         area,
-        i as usize,
-        i == eob - 1,
+        i,
+        i == usize::from(eob) - 1,
         tx_size,
         tx_class,
-      ) as i8;
+      ) as i8);
     }
+    // SAFETY: every element has been initialized
+    unsafe { slice_assume_init_mut(coeffs) }
   }
 
   pub fn get_br_ctx(