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

ai_edge_litert/source_context_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/source_context.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/source_context.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n$google/protobuf/source_context.proto\x12\x0fgoogle.protobuf\"\"\n\rSourceContext\x12\x11\n\tfile_name\x18\x01 \x01(\tB\x8a\x01\n\x13\x63om.google.protobufB\x12SourceContextProtoP\x01Z6google.golang.org/protobuf/types/known/sourcecontextpb\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.source_context_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\022SourceContextProtoP\001Z6google.golang.org/protobuf/types/known/sourcecontextpb\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_SOURCECONTEXT']._serialized_start=57
+  _globals['_SOURCECONTEXT']._serialized_end=91
+# @@protoc_insertion_point(module_scope)

ai_edge_litert/struct_pb2.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: google/protobuf/struct.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/struct.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cgoogle/protobuf/struct.proto\x12\x0fgoogle.protobuf\"\x84\x01\n\x06Struct\x12\x33\n\x06\x66ields\x18\x01 \x03(\x0b\x32#.google.protobuf.Struct.FieldsEntry\x1a\x45\n\x0b\x46ieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.google.protobuf.Value:\x02\x38\x01\"\xea\x01\n\x05Value\x12\x30\n\nnull_value\x18\x01 \x01(\x0e\x32\x1a.google.protobuf.NullValueH\x00\x12\x16\n\x0cnumber_value\x18\x02 \x01(\x01H\x00\x12\x16\n\x0cstring_value\x18\x03 \x01(\tH\x00\x12\x14\n\nbool_value\x18\x04 \x01(\x08H\x00\x12/\n\x0cstruct_value\x18\x05 \x01(\x0b\x32\x17.google.protobuf.StructH\x00\x12\x30\n\nlist_value\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.ListValueH\x00\x42\x06\n\x04kind\"3\n\tListValue\x12&\n\x06values\x18\x01 \x03(\x0b\x32\x16.google.protobuf.Value*\x1b\n\tNullValue\x12\x0e\n\nNULL_VALUE\x10\x00\x42\x7f\n\x13\x63om.google.protobufB\x0bStructProtoP\x01Z/google.golang.org/protobuf/types/known/structpb\xf8\x01\x01\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.struct_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\013StructProtoP\001Z/google.golang.org/protobuf/types/known/structpb\370\001\001\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_STRUCT_FIELDSENTRY']._loaded_options = None
+  _globals['_STRUCT_FIELDSENTRY']._serialized_options = b'8\001'
+  _globals['_NULLVALUE']._serialized_start=474
+  _globals['_NULLVALUE']._serialized_end=501
+  _globals['_STRUCT']._serialized_start=50
+  _globals['_STRUCT']._serialized_end=182
+  _globals['_STRUCT_FIELDSENTRY']._serialized_start=113
+  _globals['_STRUCT_FIELDSENTRY']._serialized_end=182
+  _globals['_VALUE']._serialized_start=185
+  _globals['_VALUE']._serialized_end=419
+  _globals['_LISTVALUE']._serialized_start=421
+  _globals['_LISTVALUE']._serialized_end=472
+# @@protoc_insertion_point(module_scope)

ai_edge_litert/tensor_buffer.py

@@ -0,0 +1,167 @@
+# 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.
+
+"""Python wrapper for LiteRT tensor buffer."""
+
+import os
+import numpy as np
+
+# pylint: disable=g-import-not-at-top
+if not os.path.splitext(__file__)[0].endswith(
+    os.path.join("ai_edge_litert", "tensor_buffer")
+):
+  # This file is part of litert package.
+  from ai_edge_litert.litert_wrapper.tensor_buffer_wrapper import _pywrap_litert_tensor_buffer_wrapper as _tb
+else:
+  # This file is part of ai_edge_litert package.
+  from ai_edge_litert import _pywrap_litert_tensor_buffer_wrapper as _tb
+# pylint: enable=g-import-not-at-top
+
+
+class TensorBuffer:
+  """Python wrapper for LiteRtTensorBuffer.
+
+  This class provides a high-level interface to the underlying C++
+  LiteRtTensorBuffer
+  implementation, allowing for creation, reading, writing, and management of
+  tensor
+  buffers in Python.
+  """
+
+  def __init__(self, capsule):
+    """Initializes a TensorBuffer with the provided PyCapsule.
+
+    Args:
+      capsule: A PyCapsule containing a pointer to a LiteRtTensorBuffer.
+    """
+    self._capsule = capsule
+
+  @staticmethod
+  def _dtype_to_str(np_dtype):
+    """Converts a numpy dtype to a string representation.
+
+    Args:
+      np_dtype: A numpy dtype (e.g., np.float32, np.int8).
+
+    Returns:
+      String representation of the data type (e.g., "float32", "int8").
+
+    Raises:
+      ValueError: If the dtype is not supported.
+    """
+    if np_dtype == np.float32:
+      return "float32"
+    elif np_dtype == np.int32:
+      return "int32"
+    elif np_dtype == np.int8:
+      return "int8"
+    else:
+      raise ValueError(f"Unsupported dtype: {np_dtype}")
+
+  @classmethod
+  def create_from_host_memory(cls, data_array):
+    """Creates a new TensorBuffer referencing existing host memory.
+
+    Args:
+      data_array: A NumPy array (e.g., np.array([[1.0, 2.0, 3.0, 4.0]],
+        dtype=np.float32)). The dtype of the array is used.
+
+    Returns:
+      A new TensorBuffer instance.
+
+    Raises:
+      ValueError: If the input is not a NumPy array or has an unsupported dtype.
+    """
+    if not isinstance(data_array, np.ndarray):
+      raise ValueError("data_array must be a NumPy array")
+
+    dtype_str = cls._dtype_to_str(data_array.dtype)
+    num_elements = data_array.size
+
+    cap = _tb.CreateTensorBufferFromHostMemory(
+        data_array, dtype_str, num_elements
+    )
+    return cls(cap)
+
+  def write(self, data_array):
+    """Writes data to this tensor buffer.
+
+    Args:
+      data_array: A NumPy array (e.g., np.array([[1.0, 2.0, 3.0, 4.0]],
+        dtype=np.float32)). The dtype of the array is used.
+
+    Example:
+      # Using NumPy array
+      test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
+      tensor_buffer.write(test_input)
+
+    Raises:
+      ValueError: If the input is not a NumPy array or has an unsupported dtype.
+    """
+    if not isinstance(data_array, np.ndarray):
+      raise ValueError("data_array must be a NumPy array")
+
+    dtype_str = self._dtype_to_str(data_array.dtype)
+    _tb.WriteTensor(self._capsule, data_array.flatten().tolist(), dtype_str)
+
+  def read(self, num_elements: int, output_dtype):
+    """Reads data from this tensor buffer.
+
+    Args:
+      num_elements: Number of elements to read.
+      output_dtype: NumPy dtype for the output (e.g., np.float32, np.int8).
+
+    Returns:
+      A NumPy array containing the tensor data.
+
+    Example:
+      # Get output as NumPy array
+      output_array = tensor_buffer.read(4, np.float32).reshape((1, 4))
+
+    Raises:
+      ValueError: If output_dtype is not a NumPy dtype or is not supported.
+    """
+    if not isinstance(output_dtype, type) or not hasattr(
+        np, output_dtype.__name__
+    ):
+      raise ValueError(f"output_dtype must be a NumPy dtype (e.g., np.float32)")
+
+    dtype_str = self._dtype_to_str(output_dtype)
+    data_list = _tb.ReadTensor(self._capsule, num_elements, dtype_str)
+    return np.array(data_list, dtype=output_dtype)
+
+  def destroy(self):
+    """Explicitly releases resources associated with this tensor buffer.
+
+    After calling this method, the tensor buffer is no longer valid for use.
+    """
+    _tb.DestroyTensorBuffer(self._capsule)
+    self._capsule = None
+
+  @property
+  def capsule(self):
+    """Returns the underlying PyCapsule for direct C++ interoperability.
+
+    When the capsule is used with compiled_model methods, the ownership remains
+    with this TensorBuffer instance. This avoids double-free when both objects
+    try to destroy the same underlying tensor buffer.
+
+    Returns:
+      The PyCapsule containing the pointer to the LiteRtTensorBuffer, with
+      ownership remaining with this TensorBuffer instance.
+    """
+    # Once destroyed, shouldn't try to access the capsule again.
+    if self._capsule is None:
+      raise ValueError("TensorBuffer has been destroyed")
+    return self._capsule

ai_edge_litert/timestamp_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/timestamp.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/timestamp.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1fgoogle/protobuf/timestamp.proto\x12\x0fgoogle.protobuf\"+\n\tTimestamp\x12\x0f\n\x07seconds\x18\x01 \x01(\x03\x12\r\n\x05nanos\x18\x02 \x01(\x05\x42\x85\x01\n\x13\x63om.google.protobufB\x0eTimestampProtoP\x01Z2google.golang.org/protobuf/types/known/timestamppb\xf8\x01\x01\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.timestamp_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\016TimestampProtoP\001Z2google.golang.org/protobuf/types/known/timestamppb\370\001\001\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_TIMESTAMP']._serialized_start=52
+  _globals['_TIMESTAMP']._serialized_end=95
+# @@protoc_insertion_point(module_scope)