chore: import upstream snapshot with attribution

tensorflow/lite/profiling/profiler_based_calibration/BUILD

@@ -0,0 +1,79 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tfl_calibration_utils",
+    srcs = ["tfl_calibration_utils.cc"],
+    hdrs = ["tfl_calibration_utils.h"],
+    deps = [
+        ":tfl_tensor_stats_profiler",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:stateful_error_reporter",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/profiling:memory_info",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+cc_library(
+    name = "tfl_tensor_stats_profiler",
+    srcs = ["tfl_tensor_stats_profiler.cc"],
+    hdrs = ["tfl_tensor_stats_profiler.h"],
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:private_common",
+    ],
+)
+
+cc_test(
+    name = "tfl_tensor_stats_profiler_test",
+    srcs = ["tfl_tensor_stats_profiler_test.cc"],
+    deps = [
+        ":tfl_tensor_stats_profiler",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:private_common",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "tfl_calibration_utils_test",
+    srcs = ["tfl_calibration_utils_test.cc"],
+    deps = [
+        ":tfl_calibration_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/profiling:memory_info",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+pybind_extension(
+    name = "_pywrap_tfl_calibration",
+    srcs = ["tfl_calibration_wrapper.cc"],
+    module_name = "_pywrap_tfl_calibration",
+    wrap_py_init = True,
+    deps = [
+        ":tfl_calibration_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:stateful_error_reporter",
+        "//tensorflow/lite/profiling:memory_info",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)

tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils.cc

@@ -0,0 +1,106 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/profiling/memory_info.h"
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler.h"
+#include "tensorflow/lite/stateful_error_reporter.h"
+
+namespace odml {
+namespace {
+
+template <typename T>
+void UpdateTensorStats(const T* tensor_buffer, int num_elements,
+                       const char* tensor_name,
+                       std::map<std::string, Range>& tensor_stats) {
+  if (num_elements <= 0) {
+    return;
+  }
+  const Eigen::Map<const Eigen::VectorX<T>> vec(tensor_buffer, num_elements);
+  const float current_min = static_cast<float>(vec.minCoeff());
+  const float current_max = static_cast<float>(vec.maxCoeff());
+  auto [it, inserted] =
+      tensor_stats.insert({tensor_name, Range{current_min, current_max}});
+  if (!inserted) {
+    it->second.min = std::min(it->second.min, current_min);
+    it->second.max = std::max(it->second.max, current_max);
+  }
+}
+}  // namespace
+
+absl::StatusOr<std::map<std::string, Range>> InvokeWithCalibration(
+    tflite::Interpreter* interpreter, int subgraph_index,
+    tflite::StatefulErrorReporter* reporter) {
+  // Use a callback to capture tensor stats during operator invocation. The
+  // results are stored in a map of tensor name to min/max stat value.
+  std::map<std::string, Range> tensor_stats;
+  auto calc_tensor_stats = [&](const TfLiteTensor* tensor) {
+    // Skip constant tensors.
+    if (tensor->allocation_type == kTfLiteMmapRo ||
+        tensor->allocation_type == kTfLitePersistentRo) {
+      return;
+    }
+    // Skip empty or unallocated tensors.
+    if (tensor->data.raw == nullptr) {
+      return;
+    }
+    if (tensor->type == kTfLiteFloat32) {
+      UpdateTensorStats<float>(tensor->data.f, tensor->bytes / sizeof(float),
+                               tensor->name, tensor_stats);
+    } else if (tensor->type == kTfLiteInt32) {
+      UpdateTensorStats<int32_t>(tensor->data.i32,
+                                 tensor->bytes / sizeof(int32_t), tensor->name,
+                                 tensor_stats);
+    }
+  };
+
+  auto profiler = std::make_unique<odml::TensorStatsProfiler>(
+      *interpreter, calc_tensor_stats);
+  interpreter->SetProfiler(profiler.get());
+
+  auto invoke_status = interpreter->subgraph(subgraph_index)->Invoke();
+
+  // Reset the profiler to avoid dangling pointer issues when interpreter is
+  // reused.
+  interpreter->SetProfiler(nullptr);
+  for (size_t i = 0; i < interpreter->subgraphs_size(); ++i) {
+    interpreter->subgraph(i)->SetProfiler(nullptr, i);
+  }
+
+  if (invoke_status != kTfLiteOk) {
+    return absl::InternalError("InvokeWithCalibration failed" +
+                               (reporter ? ": " + reporter->message() : ""));
+  }
+
+  return tensor_stats;
+}
+
+tflite::profiling::memory::MemoryUsage GetMemoryUsage() {
+  return tflite::profiling::memory::GetMemoryUsage();
+}
+}  // namespace odml

tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils.h

@@ -0,0 +1,47 @@
+/* Copyright 2026 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 TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_CALIBRATION_UTILS_H_
+#define TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_CALIBRATION_UTILS_H_
+
+#include <map>
+#include <string>
+
+#include "absl/status/statusor.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/profiling/memory_info.h"
+#include "tensorflow/lite/stateful_error_reporter.h"
+
+namespace odml {
+
+struct Range {
+  float min;
+  float max;
+};
+
+// Invokes the subgraph at `subgraph_index` using the specified `interpreter`,
+// and returns a map of min/max values for each non-constant tensor encountered
+// during execution. Currently min/max collection supports only float32 and
+// int32 tensors, all other tensors are ignored.
+absl::StatusOr<std::map<std::string, Range>> InvokeWithCalibration(
+    tflite::Interpreter* interpreter, int subgraph_index,
+    tflite::StatefulErrorReporter* reporter = nullptr);
+
+// Gets memory usage stats.
+tflite::profiling::memory::MemoryUsage GetMemoryUsage();
+
+}  // namespace odml
+
+#endif  // TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_CALIBRATION_UTILS_H_

tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils_test.cc

@@ -0,0 +1,132 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils.h"
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "absl/status/statusor.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/profiling/memory_info.h"
+
+namespace odml {
+namespace {
+
+class TflCalibrationUtilsTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    interpreter_ = std::make_unique<tflite::Interpreter>();
+  }
+
+  std::unique_ptr<tflite::Interpreter> interpreter_;
+};
+
+TEST_F(TflCalibrationUtilsTest, InvokeWithCalibrationFloat) {
+  // 1. Add and configure 2 float32 tensors.
+  interpreter_->AddTensors(2);
+  interpreter_->SetInputs({0});
+  interpreter_->SetOutputs({1});
+  interpreter_->SetTensorParametersReadWrite(
+      /*tensor_index=*/0, /*type=*/kTfLiteFloat32, /*name=*/"input",
+      /*dims=*/{1, 4}, /*quantization=*/TfLiteQuantization());
+  interpreter_->SetTensorParametersReadWrite(
+      /*tensor_index=*/1, /*type=*/kTfLiteFloat32, /*name=*/"output",
+      /*dims=*/{1, 4}, /*quantization=*/TfLiteQuantization());
+
+  // 2. Prepare input data values.
+  const float in[] = {1.0, 2.0, -1.0, 0.0};
+
+  // 3. Register op with a dummy invoker.
+  // Counter must be static because TfLiteRegistration requires a C-style
+  // function pointer (stateless lambda), which cannot capture local variables.
+  static int invoke_count;
+  invoke_count = 0;  // Reset counter before each test.
+  TfLiteRegistration reg = {
+      .invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+        invoke_count++;
+        return kTfLiteOk;
+      }};
+  ASSERT_EQ(interpreter_->AddNodeWithParameters(
+                /*inputs=*/{0}, /*outputs=*/{1}, /*init_data=*/nullptr,
+                /*init_data_size=*/0, /*builtin_data=*/nullptr,
+                /*registration=*/&reg),
+            kTfLiteOk);
+
+  // 4. Allocate the input tensors and populate data before calibration.
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  memcpy(interpreter_->tensor(0)->data.raw, in, sizeof(in));
+
+  // 5. Run calibration step and check min/max limits.
+  auto status_or_stats = InvokeWithCalibration(interpreter_.get(), 0);
+  ASSERT_TRUE(status_or_stats.ok());
+  auto stats = *status_or_stats;
+  EXPECT_EQ(invoke_count, 1);
+  EXPECT_TRUE(stats.count("input"));
+  EXPECT_EQ(stats["input"].min, -1.0f);
+  EXPECT_EQ(stats["input"].max, 2.0f);
+}
+
+TEST_F(TflCalibrationUtilsTest, InvokeWithCalibrationInt32) {
+  // 1. Add and configure 1 int32 tensor.
+  interpreter_->AddTensors(1);
+  interpreter_->SetInputs({0});
+  interpreter_->SetTensorParametersReadWrite(
+      /*tensor_index=*/0, /*type=*/kTfLiteInt32, /*name=*/"input_i32",
+      /*dims=*/{1, 4}, /*quantization=*/TfLiteQuantization());
+
+  // 2. Prepare int32 input data values.
+  const int32_t in[] = {10, 20, -10, 0};
+
+  // 3. Register op with a dummy invoker.
+  // Counter must be static because TfLiteRegistration requires a C-style
+  // function pointer (stateless lambda), which cannot capture local variables.
+  static int invoke_count;
+  invoke_count = 0;  // Reset counter before each test.
+  TfLiteRegistration reg = {
+      .invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+        invoke_count++;
+        return kTfLiteOk;
+      }};
+  ASSERT_EQ(interpreter_->AddNodeWithParameters(
+                /*inputs=*/{0}, /*outputs=*/{}, /*init_data=*/nullptr,
+                /*init_data_size=*/0, /*builtin_data=*/nullptr,
+                /*registration=*/&reg),
+            kTfLiteOk);
+
+  // 4. Allocate the int32 tensors and populate data before calibration.
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  memcpy(interpreter_->tensor(0)->data.raw, in, sizeof(in));
+
+  // 5. Run calibration step and check min/max limits.
+  auto status_or_stats = InvokeWithCalibration(interpreter_.get(), 0);
+  ASSERT_TRUE(status_or_stats.ok());
+  auto stats = *status_or_stats;
+  EXPECT_EQ(invoke_count, 1);
+  EXPECT_TRUE(stats.count("input_i32"));
+  EXPECT_EQ(stats["input_i32"].min, -10.0f);
+  EXPECT_EQ(stats["input_i32"].max, 20.0f);
+}
+
+TEST_F(TflCalibrationUtilsTest, GetMemoryUsage) {
+  tflite::profiling::memory::MemoryUsage usage = GetMemoryUsage();
+  // We can't check exact values, but they should be filled somehow.
+  EXPECT_GT(usage.total_allocated_bytes, 0);
+}
+
+}  // namespace
+}  // namespace odml

tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_wrapper.cc

@@ -0,0 +1,92 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include <Python.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/profiling/memory_info.h"
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_calibration_utils.h"
+#include "tensorflow/lite/stateful_error_reporter.h"
+
+namespace {
+namespace py = pybind11;
+using tflite::profiling::memory::MemoryUsage;
+}  // namespace
+
+PYBIND11_MODULE(_pywrap_tfl_calibration, m) {
+  py::class_<odml::Range>(m, "Range", R"pbdoc(
+Min/max range for a tensor.
+Attributes:
+    min: min value.
+    max: max value.
+)pbdoc")
+      .def_readonly("min", &odml::Range::min)
+      .def_readonly("max", &odml::Range::max);
+  py::class_<MemoryUsage>(m, "MemoryUsage", R"pbdoc(
+Memory usage statistics.
+Attributes:
+    mem_footprint_kb: Max resident set size in KB.
+    total_allocated_bytes: Total non-mmapped heap space allocated in bytes.
+    in_use_allocated_bytes: Total heap bytes in use in bytes.
+    private_footprint_bytes: Private footprint in bytes.
+)pbdoc")
+      .def_readonly("mem_footprint_kb", &MemoryUsage::mem_footprint_kb)
+      .def_readonly("total_allocated_bytes",
+                    &MemoryUsage::total_allocated_bytes)
+      .def_readonly("in_use_allocated_bytes",
+                    &MemoryUsage::in_use_allocated_bytes)
+      .def_readonly("private_footprint_bytes",
+                    &MemoryUsage::private_footprint_bytes);
+  m.def(
+      "get_memory_usage",
+      []() {
+        py::gil_scoped_release release;
+        return odml::GetMemoryUsage();
+      },
+      R"pbdoc(
+        Returns memory usage stats.
+      )pbdoc");
+  m.def(
+      "InvokeWithCalibration",
+      [](py::object interpreter_handle, int subgraph_index) {
+        auto* interpreter = reinterpret_cast<tflite::Interpreter*>(
+            interpreter_handle.cast<intptr_t>());
+        py::gil_scoped_release release;
+        auto status_or_map = odml::InvokeWithCalibration(
+            interpreter, subgraph_index,
+            static_cast<tflite::StatefulErrorReporter*>(
+                interpreter->error_reporter()));
+        if (!status_or_map.ok())
+          throw std::runtime_error(
+              std::string(status_or_map.status().message()));
+        return status_or_map.value();
+      },
+      R"pbdoc(
+        Invoke the given ``tf.lite.Interpreter`` for calibration. Assumes
+        input tensors are already set.
+        Args:
+          interpreter: The ``tf.lite:Interpreter`` to invoke.
+          subgraph_index: The subgraph index to run.
+        Returns:
+          dict: A map of tensor names to Range objects containing 'min' and 'max' values.
+      )pbdoc");
+}

tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler.cc

@@ -0,0 +1,104 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler.h"
+
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace odml {
+
+TensorStatsProfiler::TensorStatsProfiler(const tflite::Interpreter& interpreter,
+                                         Callback callback)
+    : interpreter_(interpreter), callback_(callback) {}
+
+uint32_t TensorStatsProfiler::BeginEvent(const char* tag, EventType event_type,
+                                         int64_t event_metadata1,
+                                         int64_t event_metadata2) {
+  // Process subgraph inputs at the beginning of each subgraph invoke.
+  if (tag && strcmp(tag, "Invoke") == 0) {
+    const int64_t subgraph_idx = event_metadata2;
+    const tflite::Subgraph* subgraph = interpreter_.subgraph(subgraph_idx);
+    if (subgraph) {
+      for (const int input_tensor_index : subgraph->inputs()) {
+        if (input_tensor_index != kTfLiteOptionalTensor) {
+          callback_(subgraph->tensor(input_tensor_index));
+        }
+      }
+    } else {
+      TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
+                      "TensorStatsProfiler: subgraph %d not found.",
+                      subgraph_idx);
+    }
+    return 0;  // No event handle for input tensors.
+  }
+
+  // Only capture operator invoke events for intermediate/output activations.
+  if (event_type != EventType::OPERATOR_INVOKE_EVENT) {
+    return 0;
+  }
+
+  const int64_t node_index = event_metadata1;
+  const int64_t subgraph_index = event_metadata2;
+
+  // Store the subgraph and node index for later retrieval.
+  events_.push_back({
+      .subgraph_index = subgraph_index,
+      .node_index = node_index,
+  });
+  return events_.size();
+}
+
+void TensorStatsProfiler::EndEvent(uint32_t event_handle) {
+  if (!event_handle || events_.size() < event_handle) {
+    return;
+  }
+
+  // Retrieve the node from the event handle.
+  const auto& event = events_[event_handle - 1];
+  const tflite::Subgraph* subgraph =
+      interpreter_.subgraph(event.subgraph_index);
+  if (!subgraph) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
+                    "TensorStatsProfiler: subgraph %d not found.",
+                    event.subgraph_index);
+    return;
+  }
+  const std::pair<TfLiteNode, TfLiteRegistration>* node_and_reg =
+      subgraph->node_and_registration(event.node_index);
+  if (!node_and_reg) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
+                    "TensorStatsProfiler: node %d in subgraph %d not found.",
+                    event.node_index, event.subgraph_index);
+    return;
+  }
+  const TfLiteNode& node = node_and_reg->first;
+
+  // Invoke the callback for each output tensor of the operator.
+  for (int i = 0; i < node.outputs->size; ++i) {
+    const int tensor_index = node.outputs->data[i];
+    if (tensor_index != kTfLiteOptionalTensor) {
+      callback_(subgraph->tensor(tensor_index));
+    }
+  }
+}
+
+}  // namespace odml

tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler.h

@@ -0,0 +1,70 @@
+/* Copyright 2026 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 TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_TENSOR_STATS_PROFILER_H_
+#define TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_TENSOR_STATS_PROFILER_H_
+
+#include <cstdint>
+#include <functional>
+#include <vector>
+
+#include "tensorflow/lite/core/api/profiler.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/interpreter.h"
+
+namespace odml {
+
+// The TensorStatsProfiler collects stats for each non-constant tensor in a
+// subgraph during interpreter invocation.
+class TensorStatsProfiler : public tflite::Profiler {
+ public:
+  using Callback = std::function<void(const TfLiteTensor*)>;
+
+  TensorStatsProfiler(const tflite::Interpreter& interpreter,
+                      Callback callback);
+
+  // This profiler may be invoked at multiple points throughout the
+  // execution of a subgraph. At the beginning of each subgraph invoke,
+  // capture the input tensor stats with the provided callback. At the beginning
+  // of each operator invoke, stores the subgraph and node index for later
+  // retrieval.
+  uint32_t BeginEvent(const char* tag, EventType event_type,
+                      int64_t event_metadata1,
+                      int64_t event_metadata2) override;
+
+  // At the end of an operator invoke event, calculates the tensor stats for the
+  // operator's output tensors with the provided callback.
+  void EndEvent(uint32_t event_handle) override;
+
+ private:
+  struct EventMetadata {
+    int64_t subgraph_index;
+    int64_t node_index;
+  };
+  // A mapping between event IDs and (subgraph_index, node_index).
+  std::vector<EventMetadata> events_;
+
+  // A handle to the active TFLite interpreter.
+  const tflite::Interpreter& interpreter_;
+
+  // A user provided callback to calculate tensor stats for a given tensor. The
+  // callback signature is:
+  //  void Callback(const TfLiteTensor* tensor);
+  Callback callback_;
+};
+
+}  // namespace odml
+
+#endif  // TENSORFLOW_LITE_PROFILING_PROFILER_BASED_CALIBRATION_TFL_TENSOR_STATS_PROFILER_H_

tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler_test.cc

@@ -0,0 +1,112 @@
+/* Copyright 2026 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.
+==============================================================================*/
+
+#include "tensorflow/lite/profiling/profiler_based_calibration/tfl_tensor_stats_profiler.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/api/profiler.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/interpreter.h"
+
+namespace odml {
+namespace {
+
+TEST(TensorStatsProfilerTest, IgnoreNotOperatorInvokeEvents) {
+  tflite::Interpreter interpreter;
+  TensorStatsProfiler profiler(interpreter, [](const TfLiteTensor*) {});
+
+  // Any event type other than OPERATOR_INVOKE_EVENT should be ignored.
+  EXPECT_EQ(
+      profiler.BeginEvent("Invoke", tflite::Profiler::EventType::DEFAULT, 0, 0),
+      0);
+  EXPECT_EQ(
+      profiler.BeginEvent(
+          "Invoke", tflite::Profiler::EventType::DELEGATE_OPERATOR_INVOKE_EVENT,
+          0, 0),
+      0);
+}
+
+TEST(TensorStatsProfilerTest, HandleOperatorInvokeEvent) {
+  // 1. Build a graph with 2 inputs, 1 intermediate, and 1 output tensor.
+  tflite::Interpreter interpreter;
+  interpreter.AddTensors(4);
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({3});
+
+  // 2. Define dummy operator lifecycle callbacks.
+  TfLiteRegistration reg = {
+      .init = [](TfLiteContext* context, const char* buffer,
+                 size_t length) -> void* { return nullptr; },
+      .free = [](TfLiteContext* context, void* buffer) {},
+      .prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+        return kTfLiteOk;
+      },
+      .invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+        return kTfLiteOk;
+      },
+  };
+
+  // 3. Attach Ops to sequence the tensor dependencies.
+  // Op 0: {0, 1} -> 2.
+  ASSERT_EQ(interpreter.AddNodeWithParameters(
+                /*inputs=*/{0, 1}, /*outputs=*/{2}, /*init_data=*/nullptr,
+                /*init_data_size=*/0, /*builtin_data=*/nullptr,
+                /*registration=*/&reg),
+            kTfLiteOk);
+  // Op 1: 2 -> 3.
+  ASSERT_EQ(interpreter.AddNodeWithParameters(
+                /*inputs=*/{2}, /*outputs=*/{3}, /*init_data=*/nullptr,
+                /*init_data_size=*/0, /*builtin_data=*/nullptr,
+                /*registration=*/&reg),
+            kTfLiteOk);
+
+  // 4. Allocate memory and prepare profiler to accumulate tensors.
+  ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
+  std::vector<const TfLiteTensor*> captured_tensors;
+  TensorStatsProfiler profiler(interpreter, [&](const TfLiteTensor* t) {
+    captured_tensors.push_back(t);
+  });
+
+  // 5. Trigger sequential execution events across the operators.
+  // Simulate execution of subgraph invoke boundary.
+  const uint32_t h_subgraph =
+      profiler.BeginEvent("Invoke", tflite::Profiler::EventType::DEFAULT, 0, 0);
+  profiler.EndEvent(h_subgraph);
+
+  // Simulate execution of Op 0 (node 0).
+  const uint32_t h0 = profiler.BeginEvent(
+      "OP1", tflite::Profiler::EventType::OPERATOR_INVOKE_EVENT, 0, 0);
+  profiler.EndEvent(h0);
+
+  // Simulate execution of Op 1 (node 1).
+  const uint32_t h1 = profiler.BeginEvent(
+      "OP2", tflite::Profiler::EventType::OPERATOR_INVOKE_EVENT, 1, 0);
+  profiler.EndEvent(h1);
+
+  // 6. Verify that all input, intermediate, and output tensors are captured
+  // correctly.
+  EXPECT_EQ(captured_tensors.size(), 4);
+  EXPECT_EQ(captured_tensors[0], interpreter.tensor(0));
+  EXPECT_EQ(captured_tensors[1], interpreter.tensor(1));
+  EXPECT_EQ(captured_tensors[2], interpreter.tensor(2));
+  EXPECT_EQ(captured_tensors[3], interpreter.tensor(3));
+}
+
+}  // namespace
+}  // namespace odml