ai-edge-quantizer-nightly 0.1.0.dev20250415__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/utils/calibration_utils_test.py

@@ -14,11 +14,67 @@
 # ==============================================================================
 
 from absl.testing import parameterized
+import numpy as np
+
 from tensorflow.python.platform import googletest
+from ai_edge_quantizer import quantizer
 from ai_edge_quantizer.utils import calibration_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+_RNG = np.random.default_rng(66)
+
+_CALIBRATION_DATASET = {
+    "signature_1": [{
+        "cache_0": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "positions": np.zeros(shape=(1, 100), dtype=np.int32),
+        "tokens": np.zeros(shape=(1, 100), dtype=np.int32),
+    }],
+    "signature_2": [{
+        "cache_0": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "positions": (
+            _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32)
+        ),
+        "tokens": _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32),
+    }],
+}
+
 
+def _get_quant_parameters(
+    quantized_model: bytes, signature_data: dict[str, list[str]]
+) -> list[np.ndarray]:
+  """Returns the quantization parameters from the quantized model."""
+  quant_params = []
+  tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+      quantized_model
+  )
+  for signature_key, signature_names in signature_data.items():
+    signature_runner = tfl_interpreter.get_signature_runner(signature_key)
+
+    for signature_name in signature_names:
+      input_details = signature_runner.get_input_details()
+      output_details = signature_runner.get_output_details()
+      if signature_name in input_details.keys():
+        quant_param = input_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      elif signature_name in output_details.keys():
+        output_details = signature_runner.get_output_details()
+        quant_param = output_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      else:
+        raise ValueError(
+            f"Signature name {signature_name} not found in the model."
+        )
+  return quant_params
 
-class CalibrationUtilsTest(parameterized.TestCase):
+
+class CalibrationQsvAlignmentUtilsTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
       dict(
@@ -66,12 +122,126 @@ class CalibrationUtilsTest(parameterized.TestCase):
   def test_update_tensor_qsv_min_max(self, old_qsv, new_qsv, expected_qsv):
     updated_qsv = calibration_utils.min_max_update(old_qsv, new_qsv)
     if isinstance(expected_qsv["min"], list):
-      self.assertListEqual(list(updated_qsv["min"]), expected_qsv["min"])
-      self.assertListEqual(list(updated_qsv["max"]), expected_qsv["max"])
+      self.assertEqual(list(updated_qsv["min"]), expected_qsv["min"])
+      self.assertEqual(list(updated_qsv["max"]), expected_qsv["max"])
     else:
       self.assertEqual(updated_qsv["min"], expected_qsv["min"])
       self.assertEqual(updated_qsv["max"], expected_qsv["max"])
 
+  def test_calibration_utils_init_fails(self):
+    model_path = "non_existent_model.tflite"
+    with self.assertRaisesWithPredicateMatch(
+        Exception, lambda err: f"{model_path}" in str(err)
+    ):
+      calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+
+  def test_calibration_utils_init_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    self.assertNotEmpty(calib_utils._signature_runners)
+    self.assertNotEmpty(calib_utils._same_as_input_scale_ops)
+
+  def test_search_tensor_by_signature_name_succeeds_on_unconstrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    expected_tensor_name = "PartitionedCall:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "serving_default", "add"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_search_tensor_by_signature_name_succeeds_on_constrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_slice.tflite"
+    )
+    expected_tensor_name = "slice_input_tensor:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "slice", "output_0"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_align_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without aligning the quantization parameters.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Align the quantization parameters and quantize again.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    calib_utils.align_quant_stats(qsv, signature_data)
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
+  def test_update_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without updating the `signature_1`.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Update the `signature_1` with stats from `signature_2`.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    min_val, max_val = calib_utils.align_quant_stats(  # for min and max only.
+        qsv,
+        {
+            "signature_2": ["output_1_1"],
+        },
+    )
+    calib_utils.update_quant_stats(
+        qsv,
+        {
+            "signature_1": ["output_1_1"],
+        },
+        min_val,
+        max_val,
+    )
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/constrained_ops_utils.py

@@ -0,0 +1,111 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Utils for handling operators with quantization constraints."""
+
+from ai_edge_quantizer import algorithm_manager
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_litert import schema_py_generated as schema_fb  # pylint: disable=g-direct-tensorflow-import
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+def get_constrained_op_list(
+    quant_constraint: _OpQuantConstraint, verbose: bool = False
+) -> list[str]:
+  """Constructs and returns a list of constrained operators.
+
+  This is achieved by invoking all materialization functions and extracting
+  the constraint argument, using monkey patching to redirect logic to wrapper
+  functions.
+
+  Args:
+    quant_constraint: The quantization constraint to filter operators by.
+    verbose: Flag to enable verbose output.
+
+  Returns:
+    A list containing operators with the specified constraint.
+  """
+  constrained_ops = []
+
+  def materialize_standard_op_wrapper(
+      op_info: qtyping.OpInfo,
+      *_args,
+      constraint: _OpQuantConstraint = _OpQuantConstraint.NO_CONSTRAIN,
+      **_kwargs,
+  ) -> list[qtyping.TensorTransformationParams]:
+    if constraint == quant_constraint:
+      constrained_ops.append(op_info.op_name)
+    # Return dummy values to avoid exceptions.
+    dummy_value = [qtyping.TensorTransformationParams("")] * 2
+    return dummy_value
+
+  # Dummy implementation of the `_are_weights_too_small` function to support
+  # `materialize_standard_op_wrapper` above.
+  def are_weights_too_small_wrapper(*_args, **_kwargs) -> bool:
+    return False
+
+  # Dummy implementation of the `_materialize_bias_for_fc_conv_ops` function to
+  # support `materialize_standard_op_wrapper` above.
+  def materialize_bias_for_fc_conv_ops_wrapper(*_args, **_kwargs):
+    return
+
+  # Do monkey patch to intercept the `materialize_standard_op` function to
+  # support `materialize_standard_op_wrapper` above.
+  original_materialize_standard_op = common_utils.materialize_standard_op
+  original_are_weights_too_small = common_quantize._are_weights_too_small  # pylint: disable=protected-access
+  original_materialize_bias_for_fc_conv_ops = (
+      common_quantize._materialize_bias_for_fc_conv_ops  # pylint: disable=protected-access
+  )
+  common_utils.materialize_standard_op = materialize_standard_op_wrapper
+  common_quantize._are_weights_too_small = are_weights_too_small_wrapper  # pylint: disable=protected-access
+  common_quantize._materialize_bias_for_fc_conv_ops = (  # pylint: disable=protected-access
+      materialize_bias_for_fc_conv_ops_wrapper
+  )
+  minmax_func_dict = algorithm_manager.MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT
+
+  # Loop over all available materialization functions to build up a list of
+  # ops with the given constraint.
+  for op, materialize_fn in minmax_func_dict.items():
+    # Create a dummy op info to trigger the materialization.
+    mock_op = schema_fb.OperatorT()
+    mock_op.inputs = [0]
+    mock_op.outputs = [0]
+    op_info = qtyping.OpInfo(
+        op=mock_op,
+        op_name=op,
+        subgraph_op_index=0,
+        op_quant_config=qtyping.OpQuantizationConfig(),
+    )
+    materialize_fn(
+        get_tensor_quant_params_fn=None,
+        op_info=op_info,
+        graph_info=None,
+        tensor_name_to_qsv=None,
+    )
+
+  if verbose:
+    print(f"  {quant_constraint} op list: {constrained_ops}")
+
+  # Restore the original functions.
+  common_utils.materialize_standard_op = original_materialize_standard_op
+  common_quantize._are_weights_too_small = original_are_weights_too_small  # pylint: disable=protected-access
+  common_quantize._materialize_bias_for_fc_conv_ops = (  # pylint: disable=protected-access
+      original_materialize_bias_for_fc_conv_ops
+  )
+  return constrained_ops

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -0,0 +1,50 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# 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.
+# ==============================================================================
+
+from tensorflow.python.platform import googletest
+from absl.testing import parameterized
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+class ConstrainedOpsUtilsTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="same_as_input_scale",
+          constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+          expected_num_ops=18,
+      ),
+      dict(
+          testcase_name="same_as_output_scale",
+          constraint=_OpQuantConstraint.SAME_AS_OUTPUT_SCALE,
+          expected_num_ops=7,
+      ),
+      dict(
+          testcase_name="no_constrain",
+          constraint=_OpQuantConstraint.NO_CONSTRAIN,
+          expected_num_ops=25,
+      ),
+  )
+  def test_get_constrained_op_list(self, constraint, expected_num_ops):
+    constrained_ops = constrained_ops_utils.get_constrained_op_list(constraint)
+    self.assertLen(constrained_ops, expected_num_ops)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/utils/test_utils.py

@@ -18,6 +18,7 @@
 import inspect as _inspect
 import os.path as _os_path
 import sys as _sys
+from typing import Optional, Union
 
 from absl.testing import parameterized
 
@@ -31,6 +32,40 @@ _OpName = qtyping.TFLOperationName
 _TensorQuantConfig = qtyping.TensorQuantizationConfig
 _OpQuantConfig = qtyping.OpQuantizationConfig
 _AlgorithmName = quantizer.AlgorithmName
+_Numeric = Union[int, float]
+
+
+DEFAULT_ACTIVATION_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=False,
+    granularity=qtyping.QuantGranularity.TENSORWISE,
+)
+DEFAULT_WEIGHT_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=True,
+    granularity=qtyping.QuantGranularity.CHANNELWISE,
+)
+
+
+def get_static_activation_quant_setting(
+    num_bits: int, symmetric: bool
+) -> _TensorQuantConfig:
+  return _TensorQuantConfig(
+      num_bits=num_bits,
+      symmetric=symmetric,
+      granularity=qtyping.QuantGranularity.TENSORWISE,
+  )
+
+
+def get_static_op_quant_config(
+    activation_config: _TensorQuantConfig = DEFAULT_ACTIVATION_QUANT_SETTING,
+    weight_config: _TensorQuantConfig = DEFAULT_WEIGHT_QUANT_SETTING,
+) -> _OpQuantConfig:
+  return qtyping.OpQuantizationConfig(
+      activation_tensor_config=activation_config,
+      weight_tensor_config=weight_config,
+      compute_precision=_ComputePrecision.INTEGER,
+  )
 
 
 def get_path_to_datafile(path):
@@ -64,7 +99,9 @@ class BaseOpTestCase(parameterized.TestCase):
       op_name: _OpName,
       op_config: _OpQuantConfig,
       num_validation_samples: int = 4,
+      num_calibration_samples: Optional[int] = None,
       error_metric: str = 'mse',
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
   ) -> model_validator.ComparisonResult:
     """Quantizes and validates the given model with the given configurations.
 
@@ -74,7 +111,10 @@ class BaseOpTestCase(parameterized.TestCase):
       op_name: The name of the operation to be quantized.
       op_config: The configuration for the operation to be quantized.
       num_validation_samples: The number of samples to use for validation.
+      num_calibration_samples: The number of samples to use for calibration. If
+        None then it will be set to num_validation_samples * 8.
       error_metric: The error error_metric to use for validation.
+      min_max_range: The min and max of the input range.
 
     Returns:
       The comparison result of the validation.
@@ -87,15 +127,21 @@ class BaseOpTestCase(parameterized.TestCase):
         op_config=op_config,
     )
     if quantizer_instance.need_calibration:
+      if num_calibration_samples is None:
+        num_calibration_samples = num_validation_samples * 8
       calibration_data = tfl_interpreter_utils.create_random_normal_input_data(
-          quantizer_instance.float_model, num_samples=num_validation_samples * 8
+          quantizer_instance.float_model,
+          num_samples=num_calibration_samples,
+          min_max_range=min_max_range,
       )
       calibration_result = quantizer_instance.calibrate(calibration_data)
       quantization_result = quantizer_instance.quantize(calibration_result)
     else:
       quantization_result = quantizer_instance.quantize()
     test_data = tfl_interpreter_utils.create_random_normal_input_data(
-        quantization_result.quantized_model, num_samples=num_validation_samples
+        quantization_result.quantized_model,
+        num_samples=num_validation_samples,
+        min_max_range=min_max_range,
     )
     return quantizer_instance.validate(test_data, error_metric)
 
@@ -145,6 +191,7 @@ class BaseOpTestCase(parameterized.TestCase):
       expected_model_size_reduction: float,
       weight_tolerance: float = 1e-4,
       output_tolerance: float = 1e-4,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
   ):
     """Check if the quantization is successful and the result is valid."""
     validation_result = self.quantize_and_validate(
@@ -152,6 +199,7 @@ class BaseOpTestCase(parameterized.TestCase):
         algorithm_key=algorithm_key,
         op_name=op_name,
         op_config=op_config,
+        min_max_range=min_max_range,
     )
     with self.subTest(name='ModelSizeReduction'):
       self.assert_model_size_reduction_above_min_pct(
@@ -165,3 +213,28 @@ class BaseOpTestCase(parameterized.TestCase):
       self.assert_output_errors_below_tolerance(
           validation_result, output_tolerance
       )
+
+  def assert_quantization_accuracy(
+      self,
+      algorithm_key: _AlgorithmName,
+      model_path: str,
+      op_name: _OpName,
+      op_config: _OpQuantConfig,
+      num_validation_samples: int = 4,
+      num_calibration_samples: Optional[int] = None,
+      output_tolerance: float = 1e-4,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
+  ):
+    """Checks if the output errors after quantization are within the tolerance."""
+    validation_result = self.quantize_and_validate(
+        model_path=model_path,
+        algorithm_key=algorithm_key,
+        num_validation_samples=num_validation_samples,
+        num_calibration_samples=num_calibration_samples,
+        op_name=op_name,
+        op_config=op_config,
+        min_max_range=min_max_range,
+    )
+    self.assert_output_errors_below_tolerance(
+        validation_result, output_tolerance
+    )

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -20,10 +20,10 @@ from typing import Any, Optional, Union
 import immutabledict
 import numpy as np
 
+from ai_edge_litert.tools import flatbuffer_utils
 from ai_edge_quantizer import qtyping
 from ai_edge_litert import schema_py_generated as schema  # pylint:disable=g-direct-tensorflow-import
-from tensorflow.lite.tools import flatbuffer_utils  # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.platform import gfile  # pylint: disable=g-direct-tensorflow-import
+import os # tensorflow.python.platform.gfile  # pylint: disable=g-direct-tensorflow-import
 
 _TFLOpName = qtyping.TFLOperationName
 
@@ -51,11 +51,35 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.LOGISTIC: schema.BuiltinOperator.LOGISTIC,
     _TFLOpName.SLICE: schema.BuiltinOperator.SLICE,
     _TFLOpName.SUM: schema.BuiltinOperator.SUM,
+    _TFLOpName.SELECT: schema.BuiltinOperator.SELECT,
     _TFLOpName.SELECT_V2: schema.BuiltinOperator.SELECT_V2,
     _TFLOpName.STABLEHLO_COMPOSITE: schema.BuiltinOperator.STABLEHLO_COMPOSITE,
     _TFLOpName.DYNAMIC_UPDATE_SLICE: (
         schema.BuiltinOperator.DYNAMIC_UPDATE_SLICE
     ),
+    _TFLOpName.PAD: schema.BuiltinOperator.PAD,
+    _TFLOpName.SQUARED_DIFFERENCE: schema.BuiltinOperator.SQUARED_DIFFERENCE,
+    _TFLOpName.MAX_POOL_2D: schema.BuiltinOperator.MAX_POOL_2D,
+    _TFLOpName.RESIZE_BILINEAR: schema.BuiltinOperator.RESIZE_BILINEAR,
+    _TFLOpName.RESIZE_NEAREST_NEIGHBOR: (
+        schema.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR
+    ),
+    _TFLOpName.GATHER_ND: schema.BuiltinOperator.GATHER_ND,
+    _TFLOpName.PACK: schema.BuiltinOperator.PACK,
+    _TFLOpName.UNPACK: schema.BuiltinOperator.UNPACK,
+    _TFLOpName.DIV: schema.BuiltinOperator.DIV,
+    _TFLOpName.BROADCAST_TO: schema.BuiltinOperator.BROADCAST_TO,
+    _TFLOpName.SQRT: schema.BuiltinOperator.SQRT,
+    _TFLOpName.GATHER: schema.BuiltinOperator.GATHER,
+    _TFLOpName.HARD_SWISH: schema.BuiltinOperator.HARD_SWISH,
+    _TFLOpName.MAXIMUM: schema.BuiltinOperator.MAXIMUM,
+    _TFLOpName.PADV2: schema.BuiltinOperator.PADV2,
+    _TFLOpName.REDUCE_MIN: schema.BuiltinOperator.REDUCE_MIN,
+    _TFLOpName.EQUAL: schema.BuiltinOperator.EQUAL,
+    _TFLOpName.NOT_EQUAL: schema.BuiltinOperator.NOT_EQUAL,
+    _TFLOpName.MIRROR_PAD: schema.BuiltinOperator.MIRROR_PAD,
+    _TFLOpName.SPACE_TO_DEPTH: schema.BuiltinOperator.SPACE_TO_DEPTH,
+    _TFLOpName.RELU: schema.BuiltinOperator.RELU,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(
@@ -86,7 +110,7 @@ TENSOR_TYPE_TO_CODE = immutabledict.immutabledict(
     (reversed(item) for item in TENSOR_CODE_TO_TYPE.items())
 )
 
-# Expose functions in tensorflow.lite.tools.flatbuffer_utils
+# Expose functions in litert.python.tools.flatbuffer_utils
 write_model = flatbuffer_utils.write_model
 
 
@@ -121,7 +145,7 @@ def get_model_content(tflite_path: str) -> bytes:
   Returns:
     The model bytes.
   """
-  with gfile.Open(tflite_path, "rb") as tflite_file:
+  with open(tflite_path, "rb") as tflite_file:
     return tflite_file.read()
 
 
@@ -134,7 +158,7 @@ def get_model_buffer(tflite_path: str) -> bytearray:
   Returns:
     model_buffer: the model buffer.
   """
-  with gfile.Open(tflite_path, "rb") as tflite_file:
+  with open(tflite_path, "rb") as tflite_file:
     return bytearray(tflite_file.read())
 
 
@@ -187,7 +211,7 @@ def parse_fc_bmm_conv_tensors(
   return input_tensor, weight_tensor, bias_tensor, output_tensor
 
 
-# flatbuffer_model has Any type since tensorflow.lite.tools.flatbuffer_utils
+# flatbuffer_model has Any type since litert.python.tools.flatbuffer_utils
 # is not type annotated.
 def buffer_to_tensors(flatbuffer_model: Any) -> dict[int, list[Any]]:
   """Returns a map from buffer id to tensors that use it."""
@@ -328,3 +352,12 @@ def get_op_side_effect_subgraphs(
     return [opts.decompositionSubgraphIndex]
   # Can add other nested ops here (control flow ops, etc).
   return []
+
+
+def get_op_name_by_index(
+    flatbuffer_model: Any, subgraph_id: int, op_index: int
+) -> str:
+  """Get the op name from the flatbuffer model."""
+  op = flatbuffer_model.subgraphs[subgraph_id].operators[op_index]
+  builtin_code = flatbuffer_model.operatorCodes[op.opcodeIndex].builtinCode
+  return TFL_OP_CODE_TO_NAME[builtin_code]