xmos-ai-tools 1.1.2.dev216__py3-none-macosx_11_0_arm64.whl → 1.1.2.dev236__py3-none-macosx_11_0_arm64.whl

xmos_ai_tools/runtime/include/tensorflow/lite/micro/kernels/xtensa/xtensa_depthwise_conv.h

@@ -25,9 +25,9 @@ namespace tflite {
 struct XtensaDepthwiseConvOpData {
   OpDataConv reference_op_data;
 
-#if defined(HIFI4) || defined(HIFI5)
+#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
   int scratch_tensor_index;
-#endif  // defined(HIFI4) || defined(HIFI5)
+#endif  // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
 
 #if defined(VISION_P6)
   int8_t* reorder_coefficient_bias;  // buffers used to keep reordered coeff and
@@ -39,7 +39,7 @@ struct XtensaDepthwiseConvOpData {
 #endif  // VISION_P6
 };
 
-#if defined(HIFI4) || defined(HIFI5)
+#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
 TfLiteStatus DepthwiseConvPrepareHifi(TfLiteContext* context, TfLiteNode* node);
 
 TfLiteStatus DepthwiseConvEvalHifi(TfLiteContext* context, TfLiteNode* node,
@@ -52,7 +52,7 @@ TfLiteStatus DepthwiseConvEvalHifi(TfLiteContext* context, TfLiteNode* node,
 
 TfLiteStatus DepthwiseConvReferenceEvalInt8(TfLiteContext* context,
                                             TfLiteNode* node);
-#endif  // defined(HIFI4) || defined(HIFI5)
+#endif  // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
 
 #if defined(VISION_P6)
 

xmos_ai_tools/runtime/include/tensorflow/lite/micro/kernels/xtensa/xtensa_reshape.h

@@ -1,4 +1,4 @@
-/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -19,10 +19,8 @@ limitations under the License.
 
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/types.h"
-namespace tflite {
 
-constexpr int kReshapeInputTensor = 0;
-constexpr int kReshapeOutputTensor = 0;
+namespace tflite {
 
 #if defined(VISION_P6)
 

xmos_ai_tools/runtime/include/tensorflow/lite/micro/kernels/xtensa/xtensa_softmax.h

@@ -22,12 +22,12 @@ limitations under the License.
 
 namespace tflite {
 
-#if defined(HIFI4) || defined(HIFI5)
+#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
 struct XtensaSoftmaxOpData {
   SoftmaxParams params;
   int scratch_tensor_index;
 };
-#endif  // defined(HIFI4) || defined(HIFI5)
+#endif  // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
 
 #if defined(VISION_P6)
 struct XtensaSoftmaxOpData {

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h

@@ -107,6 +107,11 @@ class GreedyMemoryPlanner : public MicroMemoryPlanner {
     return per_buffer_size;
   }
 
+  // Returns False because the GreedyMemoryPlanner doesn't preserves all tensors
+  // after invocation. Do to the fact that tensors that tensor data for tensors
+  // that aren't being used during a phase of invocation are overwritten.
+  bool preserves_all_tensors() const override { return false; }
+
   TF_LITE_REMOVE_VIRTUAL_DELETE
 
  private:

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/linear_memory_planner.h

@@ -35,6 +35,10 @@ class LinearMemoryPlanner : public MicroMemoryPlanner {
   int GetBufferCount() override;
   TfLiteStatus GetOffsetForBuffer(int buffer_index, int* offset) override;
 
+  // Returns True because the LinearMemoryPlanner preserves all tensors after
+  // invocation.
+  bool preserves_all_tensors() const override { return true; }
+
  private:
   static constexpr int kMaxBufferCount = 1024;
   size_t buffer_offsets_[kMaxBufferCount];

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/micro_memory_planner.h

@@ -81,6 +81,10 @@ class MicroMemoryPlanner {
     return kTfLiteOk;
   }
 
+  // Method will return True if the MicroMemoryPlanner preserves all tensors
+  // after invocation, and False if it doesn't.
+  virtual bool preserves_all_tensors() const = 0;
+
   virtual void PrintMemoryPlan() {
     // Default does nothing.
   }

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/non_persistent_buffer_planner_shim.h

@@ -115,6 +115,10 @@ class NonPersistentMemoryPlannerShim : public MicroMemoryPlanner {
   size_t GetMaximumMemorySize() override;
   int GetBufferCount() override;
 
+  // Returns False because the NonPersistentMemoryPlannerShim doesn't preserves
+  // all tensors after invocation.
+  bool preserves_all_tensors() const override { return false; }
+
  private:
   const BufferPlan* buffer_plan_;  // not owned, can't be null
 

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_allocator.h

@@ -18,11 +18,11 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 
-#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/micro/arena_allocator/single_arena_buffer_allocator.h"
 #include "tensorflow/lite/micro/compatibility.h"
 #include "tensorflow/lite/micro/flatbuffer_utils.h"
 #include "tensorflow/lite/micro/memory_planner/micro_memory_planner.h"
+#include "tensorflow/lite/micro/micro_common.h"
 #include "tensorflow/lite/micro/tflite_bridge/flatbuffer_conversions_bridge.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -66,9 +66,16 @@ struct ScratchBufferRequest {
 
 }  // namespace internal
 
+// Enum used to keep track of which MemoryPlanner is being used for
+// MicroAllocater::Create();
+enum class MemoryPlannerType {
+  kGreedy,
+  kLinear,
+};
+
 struct NodeAndRegistration {
   TfLiteNode node;
-  const TfLiteRegistration_V1* registration;
+  const TFLMRegistration* registration;
 };
 
 // Holds a pointer to a buffer for a scratch buffer requested by a kernel during
@@ -117,7 +124,9 @@ class MicroAllocator {
   // Note: Please use alignas(16) to make sure tensor_arena is 16
   // bytes aligned, otherwise some head room will be wasted.
   // TODO(b/157615197): Cleanup constructor + factory usage.
-  static MicroAllocator* Create(uint8_t* tensor_arena, size_t arena_size);
+  static MicroAllocator* Create(
+      uint8_t* tensor_arena, size_t arena_size,
+      MemoryPlannerType memory_planner_type = MemoryPlannerType::kGreedy);
 
   // Creates a MicroAllocator instance from a given tensor arena and a given
   // MemoryPlanner. This arena will be managed by the created instance. Note:
@@ -137,14 +146,20 @@ class MicroAllocator {
   // SingleArenaBufferAllocator instance and the MemoryPlanner. This allocator
   // instance will use the SingleArenaBufferAllocator instance to manage
   // allocations internally.
-  static MicroAllocator* Create(uint8_t* persistent_tensor_arena,
-                                size_t persistent_arena_size,
-                                uint8_t* non_persistent_tensor_arena,
-                                size_t non_persistent_arena_size);
+  static MicroAllocator* Create(
+      uint8_t* persistent_tensor_arena, size_t persistent_arena_size,
+      uint8_t* non_persistent_tensor_arena, size_t non_persistent_arena_size,
+      MemoryPlannerType memory_planner_type = MemoryPlannerType::kGreedy);
 
   // Returns the fixed amount of memory overhead of MicroAllocator.
   static size_t GetDefaultTailUsage(bool is_memory_planner_given);
 
+  // Returns True if the MicroAllocator uses a LinearMemoryPlanner(is compatible
+  // with the PerserveAllTensors flag / feature ) and False otherwise.
+  bool preserves_all_tensor() const {
+    return memory_planner_->preserves_all_tensors();
+  };
+
   // Allocates internal resources required for model inference for each subgraph
   // from the arena.
   //
@@ -185,7 +200,7 @@ class MicroAllocator {
   // Allocates a TfLiteTensor struct and populates the returned value with
   // properties from the model flatbuffer. This struct is allocated from
   // temporary arena memory is only guaranteed until a call is made to
-  // ResetTempAllocations(). Subgraph_allocaitons contains the array of
+  // ResetTempAllocations(). Subgraph_allocations contains the array of
   // TfLiteEvalTensors. If the newly allocated temp at the specified subgraph
   // and tensor index is already present int the TfLiteEvalTensor array, its
   // data buffer will be re-used.

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_common.h

@@ -0,0 +1,38 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef THIRD_PARTY_TFLITE_MICRO_TENSORFLOW_LITE_MICRO_MICRO_COMMON_H_
+#define THIRD_PARTY_TFLITE_MICRO_TENSORFLOW_LITE_MICRO_MICRO_COMMON_H_
+
+#include "tensorflow/lite/c/common.h"
+
+// TFLMRegistration defines the API that TFLM kernels need to implement.
+// This will be replacing the current TfLiteRegistration_V1 struct with
+// something more compatible Embedded enviroment TFLM is used in.
+struct TFLMRegistration {
+  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
+  void (*free)(TfLiteContext* context, void* buffer);
+  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
+  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
+  void (*reset)(TfLiteContext* context, void* buffer);
+  int32_t builtin_code;
+  const char* custom_name;
+};
+
+struct TFLMInferenceRegistration {
+  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
+  void (*reset)(TfLiteContext* context, void* buffer);
+};
+
+#endif  // THIRD_PARTY_TFLITE_MICRO_TENSORFLOW_LITE_MICRO_MICRO_COMMON_H_

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_context.h

@@ -17,7 +17,6 @@ limitations under the License.
 #define TENSORFLOW_LITE_MICRO_MICRO_CONTEXT_H_
 
 #include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/micro/micro_allocator.h"
 #include "tensorflow/lite/micro/micro_graph.h"
 
 #define XCORE_TFLITE_MICRO_PATCHED
@@ -25,6 +24,8 @@ limitations under the License.
 #ifdef NO_INTERPRETER
 
 namespace tflite {
+  const TfLiteStatus kTfLiteAbort = static_cast<TfLiteStatus>(15);
+
   struct MicroContext{
       TfLiteTensor* (*AllocateTempInputTensor)(const TfLiteNode* node, int index);
       TfLiteTensor* (*AllocateTempOutputTensor)(const TfLiteNode* node, int index);
@@ -40,6 +41,9 @@ namespace tflite {
 #else
 
 namespace tflite {
+// TODO(b/149795762): kTfLiteAbort cannot be part of the tflite TfLiteStatus.
+const TfLiteStatus kTfLiteAbort = static_cast<TfLiteStatus>(15);
+
 // MicroContext is eventually going to become the API between TFLM and the
 // kernels, replacing all the functions in TfLiteContext. The end state is code
 // kernels to have code like:
@@ -48,119 +52,73 @@ namespace tflite {
 // micro_context-><TFLM kernel API>
 class MicroContext {
  public:
-  // Enum that allows MicroContext to keep track of the stages different memory
-  // planning APIs are available to kernels.
-  enum class InterpreterState {
-    kInit,
-    kPrepare,
-    kMemoryPlanning,
-    kInvoke,
-  };
-
-  // Does not take any ownership, and all pointers must refer to valid objects
-  // that outlive the one constructed.
-  explicit MicroContext(MicroAllocator* allocator, const Model* model,
-                        MicroGraph* graph);
-  virtual ~MicroContext();
+  virtual ~MicroContext() = default;
 
   // Allocate persistent buffer which has the same life time as the interpreter.
   // Returns nullptr on failure.
   // The memory is allocated from the tail.
   // This method is only available in Init or Prepare stage.
-  // Virtual so that it can be faked for kernel tests.
-  virtual void* AllocatePersistentBuffer(size_t bytes);
+  virtual void* AllocatePersistentBuffer(size_t bytes) = 0;
 
   // Request a scratch buffer in the arena through static memory planning.
   // This method is only available in Prepare stage and the buffer is allocated
   // by the interpreter between Prepare and Eval stage. In Eval stage,
   // GetScratchBuffer API can be used to fetch the address.
-  // Virtual so that it can be faked for kernel tests.
   virtual TfLiteStatus RequestScratchBufferInArena(size_t bytes,
-                                                   int* buffer_idx);
+                                                   int* buffer_idx) = 0;
 
   // Get the scratch buffer pointer.
   // This method is only available in Eval stage.
-  // Virtual so that it can be faked for kernel tests.
-  virtual void* GetScratchBuffer(int buffer_idx);
+  virtual void* GetScratchBuffer(int buffer_idx) = 0;
 
   // Returns a temporary TfLiteTensor struct for a given index.
-  // Virtual so that it can be faked for kernel tests.
-  virtual TfLiteTensor* AllocateTempTfLiteTensor(int tensor_idx);
-  TfLiteTensor* AllocateTempTfLiteTensor(int tensor_idx, int sg);
+  virtual TfLiteTensor* AllocateTempTfLiteTensor(int tensor_idx) = 0;
+  virtual TfLiteTensor* AllocateTempTfLiteTensor(int tensor_idx, int sg){return nullptr;}
 
   // Returns a temporary TfLiteTensor struct for the specified input tensor of a
   // given mode. This is the recommended API over the deprecated
   // GetInput/GetInputSafe to get a temp input tensor. The returned tensor shall
   // be freed via calling DeallocateTempTfLiteTensor.
-  virtual TfLiteTensor* AllocateTempInputTensor(const TfLiteNode* node,
-                                                int index);
+  TfLiteTensor* AllocateTempInputTensor(const TfLiteNode* node, int index);
 
   // Returns a temporary TfLiteTensor struct for the specified output tensor of
   // a given mode. This is the recommended API over the deprecated
   // GetOutput/GetOutputSafe to get a temp output tensor. The returned tensor
   // shall be freed via calling DeallocateTempTfLiteTensor.
-  virtual TfLiteTensor* AllocateTempOutputTensor(const TfLiteNode* node,
-                                                 int index);
+  TfLiteTensor* AllocateTempOutputTensor(const TfLiteNode* node, int index);
 
   // Returns a temporary TfLiteTensor struct for the specified intermediate
   // tensor of a given mode. This is the recommended API over the deprecated
   // GetIntermediates/GetIntermediatesSafe to get a temp intermediate tensor.
   // The returned tensor shall be freed via calling DeallocateTempTfLiteTensor.
-  virtual TfLiteTensor* AllocateTempIntermediateTensor(const TfLiteNode* node,
-                                                       int index);
+  TfLiteTensor* AllocateTempIntermediateTensor(const TfLiteNode* node,
+                                               int index);
 
   // Deallocates a temp TfLiteTensor.
-  // Virtual so that it can be faked for kernel tests.
-  virtual void DeallocateTempTfLiteTensor(TfLiteTensor* tensor);
+  virtual void DeallocateTempTfLiteTensor(TfLiteTensor* tensor) = 0;
 
   // Returns a pointer to a temporary buffer (from the arena).
   // This API is only valid from the kernel's Prepare function and
   // the buffer's lifetime is also that of the Prepare function.
-  // Virtual so that it can be faked for kernel tests.
-  virtual uint8_t* AllocateTempBuffer(size_t size, size_t alignment);
+  virtual uint8_t* AllocateTempBuffer(size_t size, size_t alignment) = 0;
 
   // Signals that the temporary buffer is no longer needed.
-  // Virtual so that it can be faked for kernel tests.
-  virtual void DeallocateTempBuffer(uint8_t* buffer);
+  virtual void DeallocateTempBuffer(uint8_t* buffer) = 0;
 
   // Returns a TfLiteEvalTensor struct for a given index.
-  // Virtual so that it can be faked for kernel tests.
-  virtual TfLiteEvalTensor* GetEvalTensor(int tensor_idx);
-  TfLiteEvalTensor* GetEvalTensor(int tensor_idx, int sg);
-
-  // Sets the State of MemoryPlanning MicroContext
-  void SetInterpreterState(MicroContext::InterpreterState state);
-
-  // Sets the State of MemoryPlanning MicroContext
-  MicroContext::InterpreterState GetInterpreterState() const;
+  virtual TfLiteEvalTensor* GetEvalTensor(int tensor_idx) = 0;
+  virtual TfLiteEvalTensor* GetEvalTensor(int tensor_idx, int sg){return nullptr;}
 
   // Does not take ownership of the pointer and the pointer must refer to valid
   // an object that outlive this class instance.
   // This can only be called once to set one external context.
-  TfLiteStatus set_external_context(void* external_context_payload);
-
-  void* external_context() { return external_context_payload_; }
+  virtual TfLiteStatus set_external_context(void* external_context_payload) = 0;
 
-  MicroGraph& graph() { return graph_; }
+  virtual void* external_context() = 0;
 
-  // Sets the pointer to a list of ScratchBufferHandle instances.
-  // Not API between TFLM and kernels. Primarily used by the framework for
-  // housekeeping in MicroContext.
-  void SetScratchBufferHandles(ScratchBufferHandle* scratch_buffer_handles);
+  virtual MicroGraph& graph() = 0;
 
  private:
-  // Return the tensor index as tensor_indices[index]. tensor_indices is of
-  // max_size. Return -1 if index is not in the valid range of tensor_indices.
-  int GetTensorIndex(int index, int max_size, const int* tensor_indices);
-
-  MicroAllocator& allocator_;
-  MicroGraph& graph_;
-  const Model* model_;
-  InterpreterState state_;
-
-  ScratchBufferHandle* scratch_buffer_handles_ = nullptr;
-  void* external_context_payload_ = nullptr;
-
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };
 

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_graph.h

@@ -16,10 +16,9 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_MICRO_MICRO_GRAPH_H_
 #define TENSORFLOW_LITE_MICRO_MICRO_GRAPH_H_
 
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/micro_common.h"
 #include "tensorflow/lite/micro/micro_resource_variable.h"
-#include "tensorflow/lite/schema/schema_generated.h"
 
 #ifdef NO_INTERPRETER
 
@@ -38,79 +37,38 @@ namespace tflite {
 
 namespace tflite {
 
-// Abstracts the details of interacting with the tflite::Model.
+// Abstracts the details of interacting with the graph from the kernels
 //
-// Provides methods to access, initialize, prepare, invoke and free any
-// subgraph in the tflite::Graph.
+// Provides methods to invoke any subgraph in the tflite::Graph.
 class MicroGraph {
  public:
-  // The lifetime of the context, model, allocator and resource_variables must
-  // be at least as long as that of the graph object, since the this class may
-  // need to access them at any time. If resource_variables is a nullptr,
-  // GetResourceVariables will return a nullptr.
-  MicroGraph(TfLiteContext* context, const Model* model,
-             MicroAllocator* allocator,
-             MicroResourceVariables* resource_variables);
-  virtual ~MicroGraph();
-
-  // Sets up builtin data and calls TfLiteRegistration_V1->Init for every
-  // operator in every subgraph in the model.
-  virtual TfLiteStatus InitSubgraphs();
-
-  // Calls TfLiteRegistration_V1->Prepare for every operator in every subgraph
-  // in the model.
-  virtual TfLiteStatus PrepareSubgraphs();
-
-  // Calls TfLiteRegistration_V1->Free for every operator in every subgraph in
-  // the model.
-  virtual TfLiteStatus FreeSubgraphs();
+  virtual ~MicroGraph() = default;
 
-  // Calls TfLiteRegistration_V1->Invoke for every operator in a single subgraph
+  // Calls TFLMRegistration->Invoke for every operator in a single subgraph
   // in the model.
-  virtual TfLiteStatus InvokeSubgraph(int subgraph_idx);
-
-  // Zeros out all variable tensors in all subgraphs in the model.
-  virtual TfLiteStatus ResetVariableTensors();
+  virtual TfLiteStatus InvokeSubgraph(int subgraph_idx) = 0;
 
   // Number of tensor inputs to a specified subgraph in the model.
-  virtual size_t NumSubgraphInputs(int subgraph_idx);
+  virtual size_t NumSubgraphInputs(int subgraph_idx) = 0;
 
   // Get the specified input tensor of a specified subgraph in the model.
-  virtual TfLiteEvalTensor* GetSubgraphInput(int subgraph_idx, int input_idx);
+  virtual TfLiteEvalTensor* GetSubgraphInput(int subgraph_idx,
+                                             int input_idx) = 0;
 
   // Number of tensor outputs from a specified subgraph in the model.
-  virtual size_t NumSubgraphOutputs(int subgraph_idx);
+  virtual size_t NumSubgraphOutputs(int subgraph_idx) = 0;
 
   // Get the specified output tensor of a specified subgraph in the model.
-  virtual TfLiteEvalTensor* GetSubgraphOutput(int subgraph_idx, int output_idx);
+  virtual TfLiteEvalTensor* GetSubgraphOutput(int subgraph_idx,
+                                              int output_idx) = 0;
 
   // Number of subgraphs in the model.
-  virtual int NumSubgraphs();
-
-  // Hook to pass in subgraph allocations tracked within the interpreter,
-  // allowing MicroGraph to init / prepare / invoke subgraphs in the model.
-  void SetSubgraphAllocations(SubgraphAllocations* subgraph_allocations);
-
-  // Get the current subgraph index. Within an on operator, this is guaranteed
-  // to be the subgraph of that operator.
-  int GetCurrentSubgraphIndex() { return current_subgraph_index_; }
-
-  // Gets the list of alloctions for each subgraph. This is the source of truth
-  // for all per-subgraph allocation data.
-  SubgraphAllocations* GetAllocations() { return subgraph_allocations_; }
+  virtual int NumSubgraphs() = 0;
 
   // Get the resource variables for this TFLM graph.
-  MicroResourceVariables* GetResourceVariables() { return resource_variables_; }
+  virtual MicroResourceVariables* GetResourceVariables() = 0;
 
  private:
-  TfLiteContext* context_;
-  const Model* model_;
-  MicroAllocator* allocator_;
-  SubgraphAllocations* subgraph_allocations_ = nullptr;
-  int current_subgraph_index_;
-  MicroResourceVariables* resource_variables_;
-  const flatbuffers::Vector<flatbuffers::Offset<SubGraph>>* subgraphs_;
-
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };
 

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_interpreter.h

@@ -24,8 +24,8 @@ limitations under the License.
 #include "tensorflow/lite/core/api/error_reporter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/micro/micro_allocator.h"
-#include "tensorflow/lite/micro/micro_context.h"
-#include "tensorflow/lite/micro/micro_graph.h"
+#include "tensorflow/lite/micro/micro_interpreter_context.h"
+#include "tensorflow/lite/micro/micro_interpreter_graph.h"
 #include "tensorflow/lite/micro/micro_op_resolver.h"
 #include "tensorflow/lite/micro/micro_profiler_interface.h"
 #include "tensorflow/lite/portable_type_to_tflitetype.h"
@@ -50,7 +50,8 @@ class MicroInterpreter {
   MicroInterpreter(const Model* model, const MicroOpResolver& op_resolver,
                    uint8_t* tensor_arena, size_t tensor_arena_size,
                    MicroResourceVariables* resource_variables = nullptr,
-                   MicroProfilerInterface* profiler = nullptr);
+                   MicroProfilerInterface* profiler = nullptr,
+                   bool preserve_all_tensors = false);
 
   // Create an interpreter instance using an existing MicroAllocator instance.
   // This constructor should be used when creating an allocator that needs to
@@ -115,6 +116,9 @@ class MicroInterpreter {
     return nullptr;
   }
 
+  // Returns a pointer to the tensor for the corresponding tensor_index
+  TfLiteEvalTensor* GetTensor(int tensor_index, int subgraph_index = 0);
+
   // Reset the state to be what you would expect when the interpreter is first
   // created. i.e. after Init and Prepare is called for the very first time.
   TfLiteStatus Reset();
@@ -135,6 +139,13 @@ class MicroInterpreter {
   // arena_used_bytes() + 16.
   size_t arena_used_bytes() const { return allocator_.used_bytes(); }
 
+  // Returns True if all Tensors are being preserves
+  // TODO(b/297106074) : revisit making C++ example or test for
+  // preserve_all_tesnors
+  bool preserve_all_tensors() const {
+    return allocator_.preserves_all_tensor();
+  }
+
   size_t operators_size(int sg) const { return model_->subgraphs()->Get(sg)->operators()->size(); }
 
   // For debugging only.
@@ -158,7 +169,7 @@ class MicroInterpreter {
   const MicroOpResolver& op_resolver_;
   TfLiteContext context_ = {};
   MicroAllocator& allocator_;
-  MicroGraph graph_;
+  MicroInterpreterGraph graph_;
   bool tensors_allocated_;
 
   TfLiteStatus initialization_status_;
@@ -170,7 +181,7 @@ class MicroInterpreter {
   TfLiteTensor** input_tensors_;
   TfLiteTensor** output_tensors_;
 
-  MicroContext micro_context_;
+  MicroInterpreterContext micro_context_;
 };
 
 }  // namespace tflite