diff --git a/tensorflow/lite/micro/kernels/decode.cc b/tensorflow/lite/micro/kernels/decode.cc
index 09c490eb043..9f4d34cff15 100644
--- a/tensorflow/lite/micro/kernels/decode.cc
+++ b/tensorflow/lite/micro/kernels/decode.cc
@@ -18,12 +18,37 @@ limitations under the License.
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/micro/kernels/decode_state.h"
 #include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/micro_arena_constants.h"
 #include "tensorflow/lite/micro/micro_context.h"
 #include "tensorflow/lite/micro/micro_log.h"
 
 namespace tflite {
 namespace {
 
+TfLiteStatus SetOutputTensorData(TfLiteContext* context, const TfLiteNode* node,
+                                 size_t tensor_output_index,
+                                 TfLiteTensor* output) {
+  if (output->data.data != nullptr) {
+    // If memory has already been assigned to the tensor, leave it be
+    return kTfLiteOk;
+  }
+
+  // If alternate decompression memory is available, set the tensor data
+  // pointer now to preclude allocation by the memory planner.
+  void* alternate_decompress_mem =
+      GetMicroContext(context)->AllocateDecompressionMemory(
+          output->bytes, MicroArenaBufferAlignment());
+  if (alternate_decompress_mem != nullptr) {
+    TfLiteEvalTensor* output_eval =
+        tflite::micro::GetEvalOutput(context, node, tensor_output_index);
+    TF_LITE_ENSURE(context, output_eval != nullptr);
+    output_eval->data.data = alternate_decompress_mem;
+    output->data.data = alternate_decompress_mem;
+  }
+
+  return kTfLiteOk;
+}
+
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const size_t num_inputs = NumInputs(node);
   const size_t num_outputs = NumOutputs(node);
@@ -43,6 +68,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = nullptr;
   TfLiteStatus status = kTfLiteOk;
 
+  micro_context->ResetDecompressionMemoryAllocations();
+
   for (size_t i = 0; i < num_inputs; i += 2) {
     input = micro_context->AllocateTempInputTensor(node, i);
     if (input == nullptr) {
@@ -95,6 +122,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
         break;
     }
 
+    status = SetOutputTensorData(context, node, i / 2, output);
+    if (status != kTfLiteOk) {
+      break;
+    }
+
     if (dsp != nullptr) {
       status = dsp->Setup(*input, *ancillary, *output);
       if (status != kTfLiteOk) {
diff --git a/tensorflow/lite/micro/kernels/decode_state_huffman_test.cc b/tensorflow/lite/micro/kernels/decode_state_huffman_test.cc
index a6558a5ef7c..0030b371d14 100644
--- a/tensorflow/lite/micro/kernels/decode_state_huffman_test.cc
+++ b/tensorflow/lite/micro/kernels/decode_state_huffman_test.cc
@@ -271,7 +271,7 @@ TF_LITE_MICRO_TEST(DecodeHuffmanTable16BitsInt16Fail) {
   tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
                               outputs.size()>(
       encodes, ancillaries, outputs, expected, tflite::Register_DECODE(),
-      kTfLiteError);
+      nullptr, kTfLiteError);
 }
 
 TF_LITE_MICRO_TEST(DecodeHuffmanTable32BitsInt8) {
diff --git a/tensorflow/lite/micro/kernels/decode_state_prune_test.cc b/tensorflow/lite/micro/kernels/decode_state_prune_test.cc
index 8d87e97017b..636a5d9a746 100644
--- a/tensorflow/lite/micro/kernels/decode_state_prune_test.cc
+++ b/tensorflow/lite/micro/kernels/decode_state_prune_test.cc
@@ -575,7 +575,7 @@ TF_LITE_MICRO_TEST(DecodePruneQuantizedInvalidZeroPointInt16) {
   tflite::testing::TestDecode<kEncodes.size() + kAncillaries.size(),
                               kOutputs.size()>(
       kEncodes, kAncillaries, kOutputs, kExpected, tflite::Register_DECODE(),
-      kTfLiteError);
+      nullptr, kTfLiteError);
 }
 
 TF_LITE_MICRO_TESTS_END
diff --git a/tensorflow/lite/micro/kernels/decode_test.cc b/tensorflow/lite/micro/kernels/decode_test.cc
index 9645193bd68..b07afda1b14 100644
--- a/tensorflow/lite/micro/kernels/decode_test.cc
+++ b/tensorflow/lite/micro/kernels/decode_test.cc
@@ -196,4 +196,54 @@ TF_LITE_MICRO_TEST(DecodeTwoTensorsLUT) {
       encodes, ancillaries, outputs, expected, tflite::Register_DECODE());
 }
 
+TF_LITE_MICRO_TEST(DecodeWithAltDecompressionMemory) {
+  // Align the tensor data the same as a Buffer in the TfLite schema
+  alignas(16) int8_t output_data[std::size(kExpectLUT0)] = {};
+  alignas(16) const AncillaryData<int8_t, std::size(kAncillaryDataLUT0)>
+      kAncillaryData = {{kDcmLUT0}, {kAncillaryDataLUT0}};
+
+  constexpr int kAncillaryShapeLUT[] = {1, sizeof(kAncillaryData)};
+
+  const TfLiteIntArray* const encoded_dims =
+      tflite::testing::IntArrayFromInts(kEncodedShapeLUT);
+  static const TensorInDatum tid_encode = {
+      kEncodedLUT,
+      *encoded_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> encodes = {
+      &tid_encode,
+  };
+
+  const TfLiteIntArray* const ancillary_dims =
+      tflite::testing::IntArrayFromInts(kAncillaryShapeLUT);
+  static const TensorInDatum tid_ancillary = {
+      &kAncillaryData,
+      *ancillary_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> ancillaries = {
+      &tid_ancillary};
+
+  const TfLiteIntArray* const output_dims =
+      tflite::testing::IntArrayFromInts(kOutputShapeLUT);
+  constexpr int kOutputZeroPointsData[] = {0};
+  const TfLiteIntArray* const kOutputZeroPoints =
+      tflite::testing::IntArrayFromInts(kOutputZeroPointsData);
+  const TfLiteFloatArray kOutputScales = {kOutputZeroPoints->size};
+  static const TensorOutDatum tod = {
+      nullptr,  // using alternate decompression memory
+      *output_dims, kTfLiteInt8, kOutputScales, *kOutputZeroPoints, 0, {},
+  };
+  static constexpr std::initializer_list<const TensorOutDatum*> outputs = {
+      &tod};
+
+  const std::initializer_list<const void*> expected = {kExpectLUT0};
+
+  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion> amr = {
+      {output_data, sizeof(output_data)}};
+
+  tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
+                              outputs.size()>(
+      encodes, ancillaries, outputs, expected, tflite::Register_DECODE(), &amr);
+}
+
 TF_LITE_MICRO_TESTS_END
diff --git a/tensorflow/lite/micro/kernels/decode_test_helpers.h b/tensorflow/lite/micro/kernels/decode_test_helpers.h
index 952bb544008..c171dacaa46 100644
--- a/tensorflow/lite/micro/kernels/decode_test_helpers.h
+++ b/tensorflow/lite/micro/kernels/decode_test_helpers.h
@@ -83,7 +83,9 @@ TfLiteStatus CheckOutput(const TfLiteTensor& output,
 template <size_t kNumInputs, size_t kNumOutputs>
 TfLiteStatus ExecuteDecodeTest(
     TfLiteTensor* tensors, const TFLMRegistration& registration,
-    const std::initializer_list<const void*>& expected) {
+    const std::initializer_list<const void*>& expected,
+    const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
+        nullptr) {
   int kInputArrayData[kNumInputs + 1] = {kNumInputs};
   for (size_t i = 0; i < kNumInputs; i++) {
     kInputArrayData[i + 1] = i;
@@ -99,6 +101,10 @@ TfLiteStatus ExecuteDecodeTest(
   micro::KernelRunner runner(registration, tensors, kNumInputs + kNumOutputs,
                              inputs_array, outputs_array, nullptr);
 
+  if (amr != nullptr) {
+    runner.GetFakeMicroContext()->SetDecompressionMemory(*amr);
+  }
+
   if (runner.InitAndPrepare() != kTfLiteOk || runner.Invoke() != kTfLiteOk) {
     return kTfLiteError;
   }
@@ -135,12 +141,15 @@ TfLiteStatus ExecuteDecodeTest(
 }
 
 template <size_t kNumInputs, size_t kNumOutputs>
-void TestDecode(const std::initializer_list<const TensorInDatum*>& encodes,
-                const std::initializer_list<const TensorInDatum*>& ancillaries,
-                const std::initializer_list<const TensorOutDatum*>& outputs,
-                const std::initializer_list<const void*>& expected,
-                const TFLMRegistration& registration,
-                const TfLiteStatus expected_status = kTfLiteOk) {
+void TestDecode(
+    const std::initializer_list<const TensorInDatum*>& encodes,
+    const std::initializer_list<const TensorInDatum*>& ancillaries,
+    const std::initializer_list<const TensorOutDatum*>& outputs,
+    const std::initializer_list<const void*>& expected,
+    const TFLMRegistration& registration,
+    const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
+        nullptr,
+    const TfLiteStatus expected_status = kTfLiteOk) {
   TfLiteTensor tensors[kNumInputs + kNumOutputs] = {};
 
   for (size_t i = 0; i < kNumInputs; i += 2) {
@@ -173,7 +182,7 @@ void TestDecode(const std::initializer_list<const TensorInDatum*>& encodes,
   }
 
   TfLiteStatus s = ExecuteDecodeTest<kNumInputs, kNumOutputs>(
-      tensors, registration, expected);
+      tensors, registration, expected, amr);
   TF_LITE_MICRO_EXPECT_EQ(s, expected_status);
 }
 
diff --git a/tensorflow/lite/micro/kernels/kernel_runner.h b/tensorflow/lite/micro/kernels/kernel_runner.h
index 8dbd7f8b015..aa644d01f32 100644
--- a/tensorflow/lite/micro/kernels/kernel_runner.h
+++ b/tensorflow/lite/micro/kernels/kernel_runner.h
@@ -67,6 +67,9 @@ class KernelRunner {
   // to stub out MicroGraph methods and track invocations on each subgraph.
   MockMicroGraph* GetMockGraph() { return &mock_micro_graph_; }
 
+  // Returns a pointer to the internal FakeMicroContext.
+  FakeMicroContext* GetFakeMicroContext() { return &fake_micro_context_; }
+
   // Returns true if all temp buffer in tests are deallocated.
   // TODO(b/209453859): move this function to private after deallocation checks
   // are enabled for all kernel tests.
diff --git a/tensorflow/lite/micro/micro_context.cc b/tensorflow/lite/micro/micro_context.cc
index ea4fd8e8dc7..ffcecb7c99d 100644
--- a/tensorflow/lite/micro/micro_context.cc
+++ b/tensorflow/lite/micro/micro_context.cc
@@ -15,11 +15,13 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/micro_context.h"
 
+#include <algorithm>
 #include <cstdarg>
 #include <cstddef>
 
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/kernels/decompress.h"
+#include "tensorflow/lite/micro/memory_helpers.h"
 #include "tensorflow/lite/micro/micro_common.h"
 #include "tensorflow/lite/micro/micro_log.h"
 #include "tensorflow/lite/micro/micro_utils.h"
@@ -125,18 +127,50 @@ void* MicroContext::DecompressTensorToBuffer(
   return nullptr;
 }
 
+#endif  // USE_TFLM_COMPRESSION
+
 TfLiteStatus MicroContext::SetDecompressionMemory(
     const std::initializer_list<AlternateMemoryRegion>& regions) {
-  return kTfLiteError;
+  if (decompress_regions_ != nullptr) {
+    return kTfLiteError;
+  }
+
+  decompress_regions_ = &regions;
+  decompress_regions_allocations_ = static_cast<size_t*>(
+      AllocatePersistentBuffer(sizeof(size_t) * regions.size()));
+  if (decompress_regions_allocations_ == nullptr) {
+    return kTfLiteError;
+  }
+  ResetDecompressionMemoryAllocations();
+
+  return kTfLiteOk;
 }
 
 void* MicroContext::AllocateDecompressionMemory(size_t bytes,
                                                 size_t alignment) {
+  if (decompress_regions_ != nullptr) {
+    for (size_t i = 0; i < decompress_regions_->size(); i++) {
+      const AlternateMemoryRegion* region = &decompress_regions_->begin()[i];
+      uint8_t* start = static_cast<uint8_t*>(region->address) +
+                       decompress_regions_allocations_[i];
+      uint8_t* aligned_start = AlignPointerUp(start, alignment);
+      size_t total = bytes + (aligned_start - start);
+      if (total + decompress_regions_allocations_[i] <= region->bytes) {
+        decompress_regions_allocations_[i] += total;
+        return aligned_start;
+      }
+    }
+  }
+
   return nullptr;
 }
 
-void MicroContext::ResetDecompressionMemoryAllocations() {}
-
-#endif  // USE_TFLM_COMPRESSION
+void MicroContext::ResetDecompressionMemoryAllocations() {
+  if (decompress_regions_ == nullptr) {
+    return;
+  }
+  TFLITE_DCHECK(decompress_regions_allocations_ != nullptr);
+  std::fill_n(decompress_regions_allocations_, decompress_regions_->size(), 0);
+}
 
 }  // namespace tflite
diff --git a/tensorflow/lite/micro/micro_context.h b/tensorflow/lite/micro/micro_context.h
index 5b1ea9ca798..f9a70bf5d32 100644
--- a/tensorflow/lite/micro/micro_context.h
+++ b/tensorflow/lite/micro/micro_context.h
@@ -16,14 +16,15 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_MICRO_MICRO_CONTEXT_H_
 #define TENSORFLOW_LITE_MICRO_MICRO_CONTEXT_H_
 
+#include <cstddef>
+#include <initializer_list>
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/micro/micro_graph.h"
 #include "tensorflow/lite/micro/micro_profiler_interface.h"
 
 #ifdef USE_TFLM_COMPRESSION
 
-#include <initializer_list>
-
 #include "tensorflow/lite/micro/compression.h"
 
 #endif  // USE_TFLM_COMPRESSION
@@ -126,6 +127,8 @@ class MicroContext {
       const TfLiteEvalTensor& tensor,
       const CompressionTensorData& compression_data, void* buffer);
 
+#endif  // USE_TFLM_COMPRESSION
+
   // Used for configuring alternate decompression memory
   struct AlternateMemoryRegion {
     void* address;
@@ -140,14 +143,13 @@ class MicroContext {
   // Return a pointer to memory that can be used for decompression.
   // The pointer will be aligned to the <alignment> value.
   // Return nullptr if the requested size is not available.
-  // Can be called during kPrepare and kInvoke states.
+  // Can be called during kPrepare state.
   virtual void* AllocateDecompressionMemory(size_t bytes, size_t alignment);
 
-  // reset all allocation tracking
+  // Reset all allocation tracking.
+  // Can be called during kPrepare state.
   virtual void ResetDecompressionMemoryAllocations();
 
-#endif  // USE_TFLM_COMPRESSION
-
   // Set the alternate MicroProfilerInterface.
   // This can be used to profile subsystems simultaneously with the profiling
   // of kernels during the Eval phase.  See (b/379584353).
@@ -168,6 +170,11 @@ class MicroContext {
   }
 
  private:
+  const std::initializer_list<AlternateMemoryRegion>* decompress_regions_ =
+      nullptr;
+  // array of size_t elements with length equal to decompress_regions_.size()
+  size_t* decompress_regions_allocations_ = nullptr;
+
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };
 
diff --git a/tensorflow/lite/micro/micro_interpreter.cc b/tensorflow/lite/micro/micro_interpreter.cc
index 666516c6a3a..0de18d3a928 100644
--- a/tensorflow/lite/micro/micro_interpreter.cc
+++ b/tensorflow/lite/micro/micro_interpreter.cc
@@ -339,14 +339,10 @@ TfLiteStatus MicroInterpreter::SetAlternateProfiler(
   return micro_context_.SetAlternateProfiler(alt_profiler);
 }
 
-#ifdef USE_TFLM_COMPRESSION
-
 TfLiteStatus MicroInterpreter::SetDecompressionMemory(
     const std::initializer_list<MicroInterpreterContext::AlternateMemoryRegion>&
         regions) {
   return micro_context_.SetDecompressionMemory(regions);
 }
 
-#endif  // USE_TFLM_COMPRESSION
-
 }  // namespace tflite
diff --git a/tensorflow/lite/micro/micro_interpreter.h b/tensorflow/lite/micro/micro_interpreter.h
index 4a03c3fe825..e47c2b8ef0f 100644
--- a/tensorflow/lite/micro/micro_interpreter.h
+++ b/tensorflow/lite/micro/micro_interpreter.h
@@ -160,8 +160,6 @@ class MicroInterpreter {
   // decompression subsystem.
   TfLiteStatus SetAlternateProfiler(MicroProfilerInterface* alt_profiler);
 
-#ifdef USE_TFLM_COMPRESSION
-
   // Set the alternate decompression memory regions.
   // Can only be called during the MicroInterpreter kInit state (i.e. must
   // be called before MicroInterpreter::AllocateTensors).
@@ -169,8 +167,6 @@ class MicroInterpreter {
       const std::initializer_list<MicroContext::AlternateMemoryRegion>&
           regions);
 
-#endif  // USE_TFLM_COMPRESSION
-
  protected:
   const MicroAllocator& allocator() const { return allocator_; }
   const TfLiteContext& context() const { return context_; }
diff --git a/tensorflow/lite/micro/micro_interpreter_context.cc b/tensorflow/lite/micro/micro_interpreter_context.cc
index 62b33afe631..3bcd115729f 100644
--- a/tensorflow/lite/micro/micro_interpreter_context.cc
+++ b/tensorflow/lite/micro/micro_interpreter_context.cc
@@ -15,18 +15,12 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/micro_interpreter_context.h"
 
-#include <cstdint>
-
-#ifdef USE_TFLM_COMPRESSION
-
 #include <algorithm>
+#include <cstdint>
 
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/memory_helpers.h"
 #include "tensorflow/lite/micro/micro_arena_constants.h"
-
-#endif  // USE_TFLM_COMPRESSION
-
-#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/micro_utils.h"
 
 namespace tflite {
@@ -220,54 +214,29 @@ void* MicroInterpreterContext::DecompressTensorToBuffer(
                                                 buffer);
 }
 
+#endif  // USE_TFLM_COMPRESSION
+
 TfLiteStatus MicroInterpreterContext::SetDecompressionMemory(
     const std::initializer_list<MicroContext::AlternateMemoryRegion>& regions) {
   if (state_ != InterpreterState::kInit) {
     return kTfLiteError;
   }
 
-  decompress_regions_ = &regions;
-  decompress_regions_allocations_ = static_cast<size_t*>(
-      AllocatePersistentBuffer(sizeof(size_t) * regions.size()));
-  if (decompress_regions_allocations_ == nullptr) {
-    return kTfLiteError;
-  }
-  ResetDecompressionMemoryAllocations();
-
-  return kTfLiteOk;
+  return MicroContext::SetDecompressionMemory(regions);
 }
 
 void* MicroInterpreterContext::AllocateDecompressionMemory(size_t bytes,
                                                            size_t alignment) {
+#ifdef USE_TFLM_COMPRESSION
   TFLITE_DCHECK(state_ == InterpreterState::kPrepare ||
                 state_ == InterpreterState::kInvoke);
-  if (decompress_regions_ != nullptr) {
-    for (size_t i = 0; i < decompress_regions_->size(); i++) {
-      const AlternateMemoryRegion* region = &decompress_regions_->begin()[i];
-      uint8_t* start = static_cast<uint8_t*>(region->address) +
-                       decompress_regions_allocations_[i];
-      uint8_t* aligned_start = AlignPointerUp(start, alignment);
-      size_t total = bytes + (aligned_start - start);
-      if (total + decompress_regions_allocations_[i] <= region->bytes) {
-        decompress_regions_allocations_[i] += total;
-        return aligned_start;
-      }
-    }
-  }
-
-  return nullptr;
-}
+#else
+  TFLITE_DCHECK(state_ == InterpreterState::kPrepare);
+#endif  // USE_TFLM_COMPRESSION
 
-void MicroInterpreterContext::ResetDecompressionMemoryAllocations() {
-  if (decompress_regions_ == nullptr) {
-    return;
-  }
-  TFLITE_DCHECK(decompress_regions_allocations_ != nullptr);
-  std::fill_n(decompress_regions_allocations_, decompress_regions_->size(), 0);
+  return MicroContext::AllocateDecompressionMemory(bytes, alignment);
 }
 
-#endif  // USE_TFLM_COMPRESSION
-
 TfLiteStatus MicroInterpreterContext::SetAlternateProfiler(
     tflite::MicroProfilerInterface* alt_profiler) {
   alt_profiler_ = alt_profiler;
diff --git a/tensorflow/lite/micro/micro_interpreter_context.h b/tensorflow/lite/micro/micro_interpreter_context.h
index a3927580d51..5f17c1efac6 100644
--- a/tensorflow/lite/micro/micro_interpreter_context.h
+++ b/tensorflow/lite/micro/micro_interpreter_context.h
@@ -16,6 +16,9 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_MICRO_MICRO_INTERPRETER_CONTEXT_H_
 #define TENSORFLOW_LITE_MICRO_MICRO_INTERPRETER_CONTEXT_H_
 
+#include <cstddef>
+#include <initializer_list>
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/micro/micro_allocator.h"
 #include "tensorflow/lite/micro/micro_context.h"
@@ -128,6 +131,8 @@ class MicroInterpreterContext : public MicroContext {
                                  const CompressionTensorData& compression_data,
                                  void* buffer) override;
 
+#endif  // USE_TFLM_COMPRESSION
+
   // Set the alternate decompression memory regions.
   // Can only be called during the MicroInterpreter kInit state.
   TfLiteStatus SetDecompressionMemory(
@@ -136,14 +141,9 @@ class MicroInterpreterContext : public MicroContext {
   // Return a pointer to memory that can be used for decompression.
   // The pointer will be aligned to the <alignment> value.
   // Return nullptr if the requested size is not available.
-  // Can be called during kPrepare and kInvoke states.
+  // Can be called during kPrepare state.
   void* AllocateDecompressionMemory(size_t bytes, size_t alignment) override;
 
-  // reset all allocation tracking
-  void ResetDecompressionMemoryAllocations() override;
-
-#endif  // USE_TFLM_COMPRESSION
-
   // Set the alternate MicroProfilerInterface.
   // This can be used to profile subsystems simultaneously with the profiling
   // of kernels during the Eval phase.  See (b/379584353).
@@ -169,15 +169,6 @@ class MicroInterpreterContext : public MicroContext {
   void* external_context_payload_ = nullptr;
   MicroProfilerInterface* alt_profiler_ = nullptr;
 
-#ifdef USE_TFLM_COMPRESSION
-
-  const std::initializer_list<AlternateMemoryRegion>* decompress_regions_ =
-      nullptr;
-  // array of size_t elements with length equal to decompress_regions_.size()
-  size_t* decompress_regions_allocations_;
-
-#endif  // USE_TFLM_COMPRESSION
-
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };
 
diff --git a/tensorflow/lite/micro/micro_interpreter_context_test.cc b/tensorflow/lite/micro/micro_interpreter_context_test.cc
index fd7fb43831f..e61514910e8 100644
--- a/tensorflow/lite/micro/micro_interpreter_context_test.cc
+++ b/tensorflow/lite/micro/micro_interpreter_context_test.cc
@@ -32,7 +32,7 @@ tflite::MicroInterpreterContext CreateMicroInterpreterContext() {
   // the test need to place non-transient memories in static variables. This is
   // safe because tests are guaranteed to run serially.
   constexpr size_t kArenaSize = 1024;
-  static uint8_t tensor_arena[kArenaSize];
+  alignas(16) static uint8_t tensor_arena[kArenaSize];
 
   const tflite::Model* model = tflite::testing::GetSimpleMockModel();
   MicroAllocator* micro_allocator =
@@ -199,4 +199,109 @@ TF_LITE_MICRO_TEST(TestGetTempIntermediateTensor) {
   TF_LITE_MICRO_EXPECT_TRUE(invalid_output == nullptr);
 }
 
+TF_LITE_MICRO_TEST(TestSetDecompressionMemory) {
+  tflite::MicroInterpreterContext micro_context =
+      tflite::CreateMicroInterpreterContext();
+
+  constexpr size_t kAltMemorySize = 1;
+  alignas(16) uint8_t g_alt_memory[kAltMemorySize];
+  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion>
+      alt_memory_region = {{g_alt_memory, kAltMemorySize}};
+  TfLiteStatus status;
+
+  // fail during Prepare state
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kPrepare);
+  status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteError);
+
+  // fail during Invoke state
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kInvoke);
+  status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteError);
+
+  // succeed during Init state
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kInit);
+  status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteOk);
+
+  // fail on second Init state attempt
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kInit);
+  status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteError);
+}
+
+TF_LITE_MICRO_TEST(TestAllocateDecompressionMemory) {
+  tflite::MicroInterpreterContext micro_context =
+      tflite::CreateMicroInterpreterContext();
+
+  constexpr size_t kAltMemorySize = 30;
+  constexpr size_t kAllocateSize = 10;
+  alignas(16) uint8_t g_alt_memory[kAltMemorySize];
+  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion>
+      alt_memory_region = {{g_alt_memory, kAltMemorySize}};
+
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kInit);
+  TfLiteStatus status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteOk);
+
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kPrepare);
+
+  // allocate first 10 bytes
+  uint8_t* p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == &g_alt_memory[0]);
+
+  // allocate next 10 bytes
+  p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == &g_alt_memory[16]);
+
+  // fail next allocation
+  p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == nullptr);
+}
+
+TF_LITE_MICRO_TEST(TestResetDecompressionMemory) {
+  tflite::MicroInterpreterContext micro_context =
+      tflite::CreateMicroInterpreterContext();
+
+  constexpr size_t kAltMemorySize = 30;
+  constexpr size_t kAllocateSize = 10;
+  alignas(16) uint8_t g_alt_memory[kAltMemorySize];
+  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion>
+      alt_memory_region = {{g_alt_memory, kAltMemorySize}};
+
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kInit);
+  TfLiteStatus status = micro_context.SetDecompressionMemory(alt_memory_region);
+  TF_LITE_MICRO_EXPECT(status == kTfLiteOk);
+
+  micro_context.SetInterpreterState(
+      tflite::MicroInterpreterContext::InterpreterState::kPrepare);
+
+  // allocate first 10 bytes
+  uint8_t* p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == &g_alt_memory[0]);
+
+  // allocate next 10 bytes
+  p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == &g_alt_memory[16]);
+
+  micro_context.ResetDecompressionMemoryAllocations();
+
+  // allocate first 10 bytes again
+  p = static_cast<uint8_t*>(micro_context.AllocateDecompressionMemory(
+      kAllocateSize, tflite::MicroArenaBufferAlignment()));
+  TF_LITE_MICRO_EXPECT(p == &g_alt_memory[0]);
+}
+
 TF_LITE_MICRO_TESTS_END
diff --git a/tensorflow/lite/micro/tools/benchmarking/Makefile.inc b/tensorflow/lite/micro/tools/benchmarking/Makefile.inc
index a79420cb982..8094f2edbd6 100644
--- a/tensorflow/lite/micro/tools/benchmarking/Makefile.inc
+++ b/tensorflow/lite/micro/tools/benchmarking/Makefile.inc
@@ -20,14 +20,12 @@ endif
     $(GENERATED_SRCS_DIR)$(GENERIC_BENCHMARK_MODEL_DIR)$(GENERIC_BENCHMARK_MODEL_NAME)_model_data.h
 endif
 
-ifeq ($(ENABLE_COMPRESSION), yes)
 ifneq ($(GENERIC_BENCHMARK_ALT_MEM_ATTR),)
     CXXFLAGS += -DGENERIC_BENCHMARK_ALT_MEM_ATTR=$(GENERIC_BENCHMARK_ALT_MEM_ATTR)
 endif
 ifneq ($(GENERIC_BENCHMARK_ALT_MEM_SIZE),)
     CXXFLAGS += -DGENERIC_BENCHMARK_ALT_MEM_SIZE=$(GENERIC_BENCHMARK_ALT_MEM_SIZE)
 endif
-endif
 
 GENERIC_BENCHMARK_SRCS := \
 $(MICROLITE_BENCHMARK_ROOT_DIR)/generic_model_benchmark.cc \
diff --git a/tensorflow/lite/micro/tools/benchmarking/generic_model_benchmark.cc b/tensorflow/lite/micro/tools/benchmarking/generic_model_benchmark.cc
index 0f58219644b..704a5075bdc 100644
--- a/tensorflow/lite/micro/tools/benchmarking/generic_model_benchmark.cc
+++ b/tensorflow/lite/micro/tools/benchmarking/generic_model_benchmark.cc
@@ -70,11 +70,10 @@ limitations under the License.
         // !defined(GENERIC_BENCHMARK_ALT_MEM_ATTR)
 
 #if defined(GENERIC_BENCHMARK_ALT_MEM_SIZE) && \
-    defined(GENERIC_BENCHMARK_ALT_MEM_ATTR) && defined(USE_TFLM_COMPRESSION)
+    defined(GENERIC_BENCHMARK_ALT_MEM_ATTR)
 #define USE_ALT_DECOMPRESSION_MEM
 #endif  // defined(GENERIC_BENCHMARK_ALT_MEM_SIZE) &&
-        // defined(GENERIC_BENCHMARK_ALT_MEM_ATTR) &&
-        // defined(USE_TFLM_COMPRESSION)
+        // defined(GENERIC_BENCHMARK_ALT_MEM_ATTR)
 
 /*
  * Generic model benchmark.  Evaluates runtime performance of a provided
@@ -220,11 +219,6 @@ int Benchmark(const uint8_t* model_data, tflite::PrettyPrintType print_type) {
 
   alignas(16) static uint8_t tensor_arena[kTensorArenaSize];
 
-#ifdef USE_ALT_DECOMPRESSION_MEM
-  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion>
-      alt_memory_region = {{g_alt_memory, kAltMemorySize}};
-#endif  // USE_ALT_DECOMPRESSION_MEM
-
   uint32_t event_handle = profiler.BeginEvent("tflite::GetModel");
   const tflite::Model* model = tflite::GetModel(model_data);
   profiler.EndEvent(event_handle);
@@ -252,6 +246,8 @@ int Benchmark(const uint8_t* model_data, tflite::PrettyPrintType print_type) {
 #ifdef USE_ALT_DECOMPRESSION_MEM
   event_handle =
       profiler.BeginEvent("tflite::MicroInterpreter::SetDecompressionMemory");
+  std::initializer_list<tflite::MicroContext::AlternateMemoryRegion>
+      alt_memory_region = {{g_alt_memory, kAltMemorySize}};
   status = interpreter.SetDecompressionMemory(alt_memory_region);
   if (status != kTfLiteOk) {
     MicroPrintf("tflite::MicroInterpreter::SetDecompressionMemory failed");