ai-edge-litert-nightly 2.2.0.dev20260102__cp312-cp312-manylinux_2_27_x86_64.whl

ai_edge_litert/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.2.0.dev20260102"

ai_edge_litert/any_pb2.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: google/protobuf/any.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'google/protobuf/any.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19google/protobuf/any.proto\x12\x0fgoogle.protobuf\"&\n\x03\x41ny\x12\x10\n\x08type_url\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c\x42v\n\x13\x63om.google.protobufB\x08\x41nyProtoP\x01Z,google.golang.org/protobuf/types/known/anypb\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.any_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\010AnyProtoP\001Z,google.golang.org/protobuf/types/known/anypb\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_ANY']._serialized_start=46
+  _globals['_ANY']._serialized_end=84
+# @@protoc_insertion_point(module_scope)

ai_edge_litert/aot/ai_pack/export_lib.py

@@ -0,0 +1,300 @@
+# Copyright 2025 Google LLC.
+#
+# 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.
+"""Utility functions for exporting models to AI pack format."""
+
+import itertools
+import os
+import pathlib
+from typing import cast
+
+from ai_edge_litert.aot.core import common
+from ai_edge_litert.aot.core import types
+from ai_edge_litert.aot.vendors import fallback_backend
+from ai_edge_litert.aot.vendors.google_tensor import target as google_tensor_target
+from ai_edge_litert.aot.vendors.mediatek import mediatek_backend
+from ai_edge_litert.aot.vendors.mediatek import target as mtk_target
+from ai_edge_litert.aot.vendors.qualcomm import qualcomm_backend
+from ai_edge_litert.aot.vendors.qualcomm import target as qnn_target
+
+# TODO: b/407453529 - Add unittests.
+
+
+_DEVICE_TARGETING_CONFIGURATION = """<config:device-targeting-config
+    xmlns:config="http://schemas.android.com/apk/config">
+{device_groups}
+</config:device-targeting-config>"""
+
+_DEVICE_GROUP_TEMPLATE = """    <config:device-group name="{device_group_name}">
+{device_selectors}
+    </config:device-group>"""
+
+_DEVICE_SELECTOR_TEMPLATE = """        <config:device-selector>
+            <config:system-on-chip manufacturer="{soc_man}" model="{soc_model}"/>
+        </config:device-selector>"""
+
+
+def _is_mobile_device_backend(backend: types.Backend):
+  target = backend.target
+  if backend.id() == qualcomm_backend.QualcommBackend.id():
+    target = cast(qnn_target.Target, target)
+    # Non Android QNN targets.
+    if target.soc_model in (
+        qnn_target.SocModel.SA8255,
+        qnn_target.SocModel.SA8295,
+    ):
+      return False
+  return True
+
+
+def _export_model_files_to_ai_pack(
+    compiled_models: types.CompilationResult,
+    ai_pack_dir: pathlib.Path,
+    ai_pack_name: str,
+    litert_model_name: str,
+    *,
+    separate_mtk_ai_pack: bool = True,
+):
+  """Exports the model tflite files to the AI pack directory structure.
+
+  Args:
+    compiled_models: The compiled models to export.
+    ai_pack_dir: The directory to export the AI pack to.
+    ai_pack_name: The name of the AI pack.
+    litert_model_name: The name of the model in the litert format.
+    separate_mtk_ai_pack: Whether to separate the MTK AI pack. If True, the main
+      AI pack will use the fallback model for MTK targets. The MTK AI pack will
+      contain all MTK models, and empty directories for non-MTK targets.
+  """
+  fallback_model = None
+  for backend, model in compiled_models.models_with_backend:
+    if backend.target_id == fallback_backend.FallbackBackend.id():
+      fallback_model = model
+  assert fallback_model is not None, 'Fallback model is required.'
+
+  model_export_dir = ai_pack_dir / ai_pack_name / 'src/main/assets'
+  os.makedirs(model_export_dir, exist_ok=True)
+  for backend, model in compiled_models.models_with_backend:
+    if not _is_mobile_device_backend(backend):
+      continue
+    target_id = backend.target_id
+    backend_id = backend.id()
+    if backend_id == fallback_backend.FallbackBackend.id():
+      target_id = 'other'
+    elif backend_id == mediatek_backend.MediaTekBackend.id():
+      target_id = backend.target_id.replace(
+          mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
+      )
+    group_name = 'model#group_' + target_id
+    export_dir = model_export_dir / group_name
+    os.makedirs(export_dir, exist_ok=True)
+    model_export_path = export_dir / (litert_model_name + common.DOT_TFLITE)
+    if (
+        separate_mtk_ai_pack
+        and backend_id == mediatek_backend.MediaTekBackend.id()
+    ):
+      # Use the fallback model for MTK targets in main AI pack.
+      model_to_export = fallback_model
+    else:
+      model_to_export = model
+    if not model_to_export.in_memory:
+      model_to_export.load()
+    model_to_export.save(model_export_path, export_only=True)
+
+  if separate_mtk_ai_pack:
+    _export_model_files_to_mtk_ai_pack(
+        compiled_models=compiled_models,
+        ai_pack_dir=ai_pack_dir,
+        ai_pack_name=ai_pack_name + '_mtk',
+        litert_model_name=litert_model_name + '_mtk',
+    )
+
+
+def _export_model_files_to_mtk_ai_pack(
+    compiled_models: types.CompilationResult,
+    ai_pack_dir: pathlib.Path,
+    ai_pack_name: str,
+    litert_model_name: str,
+):
+  """Exports the model tflite files to the MTK AI pack directory structure."""
+  model_export_dir = ai_pack_dir / ai_pack_name / 'src/main/assets'
+  os.makedirs(model_export_dir, exist_ok=True)
+  for backend, model in compiled_models.models_with_backend:
+    if not _is_mobile_device_backend(backend):
+      continue
+    backend_id = backend.id()
+    target_id = backend.target_id
+    if backend_id == fallback_backend.FallbackBackend.id():
+      target_id = 'other'
+    elif backend_id == mediatek_backend.MediaTekBackend.id():
+      target_id = backend.target_id.replace(
+          mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
+      )
+    group_name = 'model#group_' + target_id
+    export_dir = model_export_dir / group_name
+    os.makedirs(export_dir, exist_ok=True)
+    if backend_id != mediatek_backend.MediaTekBackend.id():
+      # Skip non-MTK targets, just create a placeholder file.
+      placeholder_file = export_dir / 'placeholder.txt'
+      placeholder_file.touch()
+      continue
+    model_export_path = export_dir / (litert_model_name + common.DOT_TFLITE)
+    if not model.in_memory:
+      model.load()
+    model.save(model_export_path, export_only=True)
+
+
+def _build_targeting_config(compiled_backends: list[types.Backend]) -> str:
+  """Builds device-targeting-config in device_targeting_configuration.xml."""
+  device_groups = []
+  for backend in compiled_backends:
+    if not _is_mobile_device_backend(backend):
+      continue
+    target = backend.target
+    device_group = _target_to_ai_pack_info(target)
+    if device_group:
+      device_groups.append(device_group)
+  device_groups = '\n'.join(device_groups)
+  return _DEVICE_TARGETING_CONFIGURATION.format(device_groups=device_groups)
+
+
+def _target_to_ai_pack_info(target: types.Target) -> str | None:
+  """Builds the device group used in device_targeting_configuration.xml."""
+  if isinstance(target, qnn_target.Target):
+    group_name = str(target)
+    selector = _process_qnn_target(target)
+    device_selectors = [
+        _DEVICE_SELECTOR_TEMPLATE.format(soc_man=man, soc_model=model)
+        for man, model in selector
+    ]
+    device_selectors = '\n'.join(device_selectors)
+    device_group = _DEVICE_GROUP_TEMPLATE.format(
+        device_group_name=group_name, device_selectors=device_selectors
+    )
+    return device_group
+  elif isinstance(target, mtk_target.Target):
+    group_name = str(target).replace(
+        mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
+    )
+    # TODO: b/407453529 - Support MTK SDK Version / OS version in selector.
+    selector = _process_mtk_target(target)
+    device_selector = _DEVICE_SELECTOR_TEMPLATE.format(
+        soc_man=selector[0], soc_model=selector[1]
+    )
+    device_group = _DEVICE_GROUP_TEMPLATE.format(
+        device_group_name=group_name, device_selectors=device_selector
+    )
+    return device_group
+  elif isinstance(target, google_tensor_target.Target):
+    group_name = str(target)
+    soc_manufacturer, soc_model = _process_google_tensor_target(target)
+    device_selector = _DEVICE_SELECTOR_TEMPLATE.format(
+        soc_man=soc_manufacturer, soc_model=soc_model
+    )
+    device_group = _DEVICE_GROUP_TEMPLATE.format(
+        device_group_name=group_name, device_selectors=device_selector
+    )
+    return device_group
+  elif isinstance(target, fallback_backend.FallbackTarget):
+    # Don't need to have device selector for fallback target.
+    return None
+  else:
+    print('unsupported target ', target)
+    return None
+
+
+# TODO: b/407453529 - Auto-generate this function from CSVs.
+def _process_qnn_target(target: qnn_target.Target) -> list[tuple[str, str]]:
+  """Returns the list of (manufacturer, model) for the given QNN target."""
+  # Play cannot distinguish between Qualcomm and QTI for now.
+  manufacturer = ['Qualcomm', 'QTI']
+  models = [str(target.soc_model)]
+  return list(itertools.product(manufacturer, models))
+
+
+# TODO: b/407453529 - Auto-generate this function from CSVs.
+def _process_mtk_target(
+    target: mtk_target.Target,
+) -> tuple[str, str]:
+  """Returns tuple of (manufacturer, model) for the given MTK target."""
+  # Play cannot distinguish between Qualcomm and QTI for now.
+  return str(target.soc_manufacturer).replace(
+      mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
+  ), str(target.soc_model)
+
+
+# TODO: b/407453529 - Auto-generate this function from CSVs.
+def _process_google_tensor_target(
+    target: google_tensor_target.Target,
+) -> tuple[str, str]:
+  """Returns tuple of (manufacturer, model) for the given Google Tensor target."""
+  return str(target.soc_manufacturer), str(target.soc_model).replace('_', ' ')
+
+
+def _write_targeting_config(
+    compiled_models: types.CompilationResult, ai_pack_dir: pathlib.Path
+) -> None:
+  """Writes device_targeting_configuration.xml for the given compiled models."""
+  compiled_backends = [x for x, _ in compiled_models.models_with_backend]
+  targeting_config = _build_targeting_config(
+      compiled_backends=compiled_backends
+  )
+
+  targeting_config_path = ai_pack_dir / 'device_targeting_configuration.xml'
+  targeting_config_path.write_text(targeting_config)
+
+
+def export(
+    compiled_models: types.CompilationResult,
+    ai_pack_dir: pathlib.Path | str,
+    ai_pack_name: str,
+    litert_model_name: str,
+) -> None:
+  """Exports the compiled models to AI pack format.
+
+  This function will export the compiled models to corresponding directory
+  structure:
+
+  {ai_pack_dir}/
+    AiPackManifest.xml
+    device_targeting_configuration.xml
+    {ai_pack_name}/src/main/assets/
+      model#group_target_1/
+        {litert_model_name}.tflite
+      model#group_target_2/
+        {litert_model_name}.tflite
+      model#group_target_3/
+        {litert_model_name}.tflite
+      model#group_other/
+        {litert_model_name}.tflite
+
+  Args:
+    compiled_models: The compiled models to export.
+    ai_pack_dir: The directory to export the AI pack to.
+    ai_pack_name: The name of the AI pack.
+    litert_model_name: The name of the model in the litert format.
+  """
+  if isinstance(ai_pack_dir, str):
+    ai_pack_dir = pathlib.Path(ai_pack_dir)
+
+  ai_pack_dir.mkdir(parents=True, exist_ok=True)
+
+  _export_model_files_to_ai_pack(
+      compiled_models=compiled_models,
+      ai_pack_dir=ai_pack_dir,
+      ai_pack_name=ai_pack_name,
+      litert_model_name=litert_model_name,
+  )
+  _write_targeting_config(
+      compiled_models=compiled_models, ai_pack_dir=ai_pack_dir
+  )

ai_edge_litert/aot/aot_compile.py

@@ -0,0 +1,153 @@
+# Copyright 2025 The LiteRT 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.
+# ==============================================================================
+
+"""AOT Compilation for LiteRT model."""
+import pathlib
+import tempfile
+
+from ai_edge_litert.aot import prepare_for_npu as core
+from ai_edge_litert.aot.core import apply_plugin
+from ai_edge_litert.aot.core import components
+from ai_edge_litert.aot.core import mlir_transforms
+from ai_edge_litert.aot.core import types
+from ai_edge_litert.aot.vendors import import_vendor
+
+
+def aot_compile(
+    input_model: types.Model | str,
+    output_dir: str | pathlib.Path | None = None,
+    target: types.Target | list[types.Target] | None = None,
+    config: (
+        types.CompilationConfig | list[types.CompilationConfig] | None
+    ) = None,
+    quantizer: components.AieQuantizerT | None = None,
+    keep_going: bool = True,
+    subgraphs_to_compile: list[int] | None = None,
+    **kwargs,
+) -> types.CompilationResult:
+  """Prepares a TFLite model for NPU execution.
+
+  High level command that erforms various backend specific pre-processing steps
+  and then applies an NPU compiler to the given model.
+
+  Args:
+    input_model: The input model to compile.
+    output_dir: Directory to write the output files to. If not specified, the
+      output files will be written to the same directory as the input file.
+    target: The target to compile for. If not specified, will compile to all
+      registered targets.
+    config: The compilation config(s). Cannot be specified with target.
+    quantizer: The quantizer to use for quantization.
+    keep_going: Whether to keep going if some backends fail. If False, fail fast
+      on the first error and raise an exception.
+    subgraphs_to_compile: The subgraph index list to compile to NPU. If None,
+      compile all subgraphs.
+    **kwargs: Additional arguments to pass to the backend.
+
+  Returns:
+    Compiled models.
+  """
+  # Only one of target or config is needed.
+  if target and config:
+    raise ValueError("Cannot specify both target and config.")
+
+  if config is None:
+    if target is None:
+      target = import_vendor.AllRegisteredTarget()
+    if isinstance(target, types.Target):
+      config = types.CompilationConfig(target=target)
+    elif isinstance(target, list):
+      config = [types.CompilationConfig(target=t) for t in target]
+    else:
+      raise ValueError("Unsupported target type.")
+
+  if isinstance(input_model, str):
+    input_path = pathlib.Path(input_model)
+    input_model = types.Model.create_from_path(input_path)
+
+  # Resolve output paths.
+  temp_dir = None
+  if not output_dir:
+    if input_model.in_memory:
+      # Use a temp dir for in-memory models.
+      # The temp dir will be cleaned up after the models are compiled and loaded
+      # back to memory (i.e. function returns).
+      temp_dir = tempfile.TemporaryDirectory()
+      output_dir = temp_dir.name
+    else:
+      input_path = input_model.path
+      output_dir = input_path.parent / "_compiled_models"
+      output_dir.mkdir(parents=True, exist_ok=True)
+      output_dir = str(output_dir)
+  output_dir_path = pathlib.Path(output_dir)
+  output_dir_path.mkdir(parents=True, exist_ok=True)
+
+  if isinstance(config, types.CompilationConfig) or not config:
+    if config:
+      # Make pytype happy.
+      assert isinstance(config, types.CompilationConfig)
+      kw_config = config.to_dict() | kwargs
+    else:
+      kw_config = kwargs
+
+    backend_class = core.resolve_backend(kw_config)
+
+    quant_recipe = kw_config.get("quantize_recipe", None)
+    if quant_recipe:
+      assert quantizer is not None, "Quantizer is required for quantization."
+
+    results = core.prepare_for_npu(
+        input_model,
+        output_dir_path,
+        backend_class,
+        kw_config,
+        transforms=mlir_transforms.MlirTransforms(),
+        quantizer=quantizer,
+        plugin=apply_plugin.ApplyPlugin(
+            experimental_capture_stderr=True,
+            subgraphs_to_compile=subgraphs_to_compile,
+        ),
+        keep_going=keep_going,
+    )
+  elif isinstance(config, list):
+    kw_configs = [c.to_dict() | kwargs for c in config]
+
+    configs_with_backend = [(core.resolve_backend(c), c) for c in kw_configs]
+    requires_quantizer = any("quantize_recipe" in c for c in kw_configs)
+    if requires_quantizer and quantizer is None:
+      raise ValueError("Quantizer is required for quantization.")
+
+    results = core.prepare_for_npu_multiple_configs(
+        input_model,
+        output_dir_path,
+        configs_with_backend,
+        transforms=mlir_transforms.MlirTransforms(),
+        quantizer=quantizer,
+        plugin=apply_plugin.ApplyPlugin(
+            experimental_capture_stderr=True,
+            subgraphs_to_compile=subgraphs_to_compile,
+        ),
+        keep_going=keep_going,
+    )
+  else:
+    # Should not reach here.
+    raise ValueError("Unsupported config type.")
+
+  if temp_dir:
+    # Load the models to memory before cleaning up the temp dir.
+    results.load()
+    temp_dir.cleanup()
+
+  return results

ai_edge_litert/aot/core/apply_plugin.py

@@ -0,0 +1,148 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Wrapper for calling the apply plugin tooling."""
+
+
+import os
+import pathlib
+import re
+import subprocess
+import tempfile
+
+from ai_edge_litert.aot.core import common
+from ai_edge_litert.aot.core import components
+from ai_edge_litert.aot.core import types
+
+_BINARY = pathlib.Path("tools/apply_plugin_main")
+
+_RE_PARTITION_STATS = re.compile(
+    r"Partitioned subgraph<(\d+)>, selected (\d+) ops, from a total of "
+    r"(\d+) ops. resulted in (\d+) partitions."
+)
+
+
+class ApplyPlugin(components.ApplyPluginT):
+  """Wrapper for calling the apply plugin tooling."""
+
+  def __init__(
+      self,
+      experimental_capture_stderr: bool = False,
+      subgraphs_to_compile: list[int] | None = None,
+  ):
+    self._experimental_capture_stderr = experimental_capture_stderr
+    self._subgraphs_to_compile = subgraphs_to_compile
+
+  @property
+  def default_err(self) -> str:
+    # NOTE: Capture stderr from underlying binary.
+    return "--"
+
+  @property
+  def component_name(self) -> str:
+    return "apply_plugin"
+
+  def __call__(
+      self,
+      input_model: types.Model,
+      output_model: types.Model,
+      soc_manufacturer: str,
+      soc_model: str,
+      sdk_libs_path: str | None = None,
+      **kwargs,
+  ):
+    """Applies a plugin to the input model.
+
+    Args:
+      input_model: The path to the input model.
+      output_model: The path to the output model.
+      soc_manufacturer: The SOC manufacturer of the plugin.
+      soc_model: The SOC model of the plugin.
+      sdk_libs_path: The path to the SDK libs. If not provided,
+        the default SDK path will be used.
+      **kwargs: Additional arguments to pass to the underlying binary.
+
+    Returns:
+      The output model.
+
+    Raises:
+      ValueError: If no tflite model was created by the underying binary.
+    """
+    if input_model.in_memory:
+      tmp_file = tempfile.NamedTemporaryFile(mode="wb")
+      input_model.save(tmp_file.name)
+    else:
+      tmp_file = None
+
+    binary = common.get_resource(_BINARY)
+    args = [
+        str(binary),
+        "--cmd=apply",
+        f"--model={str(input_model.path)}",
+        f"--o={str(output_model.path)}",
+        f"--soc_manufacturer={soc_manufacturer}",
+        f"--soc_model={soc_model}",
+        f"--err={self.default_err}",
+    ]
+    extra_args = [f"--{key}={value}" for key, value in kwargs.items()]
+    args.extend(extra_args)
+    if self._subgraphs_to_compile:
+      subgraphs_to_compile = ",".join(
+          str(s) for s in self._subgraphs_to_compile
+      )
+      args.append(f"--subgraphs={subgraphs_to_compile}")
+    env = os.environ.copy()
+    ld_library_path = common.construct_ld_library_path()
+    if ld_library_path:
+      if sdk_libs_path:
+        ld_library_path = f"{sdk_libs_path}{os.pathsep}{ld_library_path}"
+      env["LD_LIBRARY_PATH"] = ld_library_path
+
+    result = subprocess.run(
+        args,
+        check=False,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        env=env,
+    )
+    if result.returncode:
+      log_file = tempfile.NamedTemporaryFile(
+          suffix=".error", mode="w", delete=False
+      )
+      log_file.write(result.stdout)
+      log_file.close()
+      raise ValueError(
+          f"{self.component_name} failed to apply plugin. See"
+          f" {log_file.name} for details."
+      )
+
+    if not common.is_tflite(output_model.path):
+      raise ValueError(f"{output_model.path} is not a TFLite model.")
+
+    partition_stats = _RE_PARTITION_STATS.findall(result.stdout)
+    output_model.partition_stats = types.PartitionStats(
+        subgraph_stats=[
+            types.SubgraphPartitionStats(
+                subgraph_index=int(s[0]),
+                num_ops_offloaded=int(s[1]),
+                num_total_ops=int(s[2]),
+                num_partitions_offloaded=int(s[3]),
+            )
+            for s in partition_stats
+        ]
+    )
+    if tmp_file is not None:
+      tmp_file.close()