# 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.

# ==============================================================================

"""TensorFlow Lite tooling helper functionality."""

import enum

import functools

import pprint

import shutil

import sys

import tempfile

import time

import warnings

from absl import logging

from google.protobuf import text_format as _text_format

from google.protobuf.message import DecodeError

from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as qc

from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as rd

from tensorflow.core.framework import graph_pb2 as _graph_pb2

from tensorflow.lite.experimental.microfrontend.python.ops import audio_microfrontend_op # pylint: disable=unused-import

# The following imports are needed to make the model_runtime_info_pb2

# and profiling_info_pb2 protos available via the litert PIP package.

from tensorflow.lite.profiling.proto import model_runtime_info_pb2 # pylint: disable=unused-import

from tensorflow.lite.profiling.proto import profiling_info_pb2 # pylint: disable=unused-import

from tensorflow.lite.python import conversion_metadata_schema_py_generated as conversion_metadata_fb

from tensorflow.lite.python import lite_constants as constants

from tensorflow.lite.python.convert import build_conversion_flags as _build_conversion_flags

from tensorflow.lite.python.convert import convert_graphdef as _convert_graphdef

from tensorflow.lite.python.convert import convert_graphdef_with_arrays as _convert_graphdef_with_arrays

from tensorflow.lite.python.convert import convert_jax_hlo as _convert_jax_hlo

from tensorflow.lite.python.convert import convert_saved_model as _convert_saved_model

from tensorflow.lite.python.convert import ConverterError # pylint: disable=unused-import

from tensorflow.lite.python.convert import deduplicate_readonly_buffers as _deduplicate_readonly_buffers

from tensorflow.lite.python.convert import mlir_quantize as _mlir_quantize

from tensorflow.lite.python.convert import mlir_sparsify as _mlir_sparsify

from tensorflow.lite.python.convert import OpsSet

from tensorflow.lite.python.convert import toco_convert # pylint: disable=unused-import

from tensorflow.lite.python.convert_phase import Component

from tensorflow.lite.python.convert_phase import convert_phase

from tensorflow.lite.python.convert_phase import SubComponent

from tensorflow.lite.python.convert_saved_model import freeze_saved_model as _freeze_saved_model

from tensorflow.lite.python.interpreter import Interpreter # pylint: disable=unused-import

from tensorflow.lite.python.interpreter import load_delegate # pylint: disable=unused-import

from tensorflow.lite.python.interpreter import OpResolverType # pylint: disable=unused-import

from tensorflow.lite.python.metrics import metrics

from tensorflow.lite.python.op_hint import convert_op_hints_to_stubs # pylint: disable=unused-import

from tensorflow.lite.python.op_hint import is_ophint_converted as _is_ophint_converted

from tensorflow.lite.python.op_hint import OpHint # pylint: disable=unused-import

from tensorflow.lite.python.optimize import calibrator as _calibrator

from tensorflow.lite.python.util import _jit

from tensorflow.lite.python.util import build_debug_info_func as _build_debug_info_func

from tensorflow.lite.python.util import convert_debug_info_func as _convert_debug_info_func

from tensorflow.lite.python.util import freeze_graph as _freeze_graph

from tensorflow.lite.python.util import get_debug_info as _get_debug_info

from tensorflow.lite.python.util import get_grappler_config as _get_grappler_config

from tensorflow.lite.python.util import get_model_hash as _get_model_hash

from tensorflow.lite.python.util import get_save_spec as _get_save_spec

from tensorflow.lite.python.util import get_sparsity_modes as _get_sparsity_modes

from tensorflow.lite.python.util import get_tensor_name as _get_tensor_name

from tensorflow.lite.python.util import get_tensors_from_tensor_names as _get_tensors_from_tensor_names

from tensorflow.lite.python.util import get_tf_type_name as _get_tf_type_name

from tensorflow.lite.python.util import is_frozen_graph as _is_frozen_graph

from tensorflow.lite.python.util import model_input_signature as _model_input_signature

from tensorflow.lite.python.util import modify_model_io_type as _modify_model_io_type

from tensorflow.lite.python.util import populate_conversion_metadata as _populate_conversion_metadata

from tensorflow.lite.python.util import run_graph_optimizations as _run_graph_optimizations

from tensorflow.lite.python.util import set_tensor_shapes as _set_tensor_shapes

from tensorflow.lite.python.util import trace_model_call as _trace_model_call

from tensorflow.lite.tools import flatbuffer_utils

from tensorflow.lite.tools.optimize.debugging.python.debugger import QuantizationDebugger # pylint: disable=unused-import

from tensorflow.lite.tools.optimize.debugging.python.debugger import QuantizationDebugOptions # pylint: disable=unused-import

from tensorflow.python.client import session as _session

from tensorflow.python.eager import context

from tensorflow.python.eager import def_function as _def_function

from tensorflow.python.eager import function as _function

from tensorflow.python.framework import byte_swap_tensor as bst

from tensorflow.python.framework import convert_to_constants as _convert_to_constants

from tensorflow.python.framework import dtypes as _dtypes

from tensorflow.python.framework import ops as _ops

from tensorflow.python.framework import tensor_spec

from tensorflow.python.framework import versions

from tensorflow.python.framework.errors_impl import NotFoundError as _NotFoundError

from tensorflow.python.framework.importer import import_graph_def as _import_graph_def

from tensorflow.python.platform import gfile

from tensorflow.python.saved_model import loader_impl as _loader_impl

from tensorflow.python.saved_model import save as _save

from tensorflow.python.saved_model import save_options as _save_options

from tensorflow.python.saved_model import signature_constants as _signature_constants

from tensorflow.python.saved_model import tag_constants as _tag_constants

from tensorflow.python.saved_model.load import load as _load

from tensorflow.python.saved_model.loader_impl import parse_saved_model_with_debug_info as _parse_saved_model_with_debug_info

from tensorflow.python.util import deprecation as _deprecation

from tensorflow.python.util import keras_deps

from tensorflow.python.util import nest

from tensorflow.python.util.tf_export import tf_export as _tf_export

@_tf_export("lite.Optimize")

class Optimize(enum.Enum):

"""Enum defining the optimizations to apply when generating a tflite model.

DEFAULT

The default optimization strategy that enables post-training quantization.

The type of post-training quantization that will be used is dependent on

the other converter options supplied. Refer to the

[documentation](https://ai.google.dev/edge/litert/models/post_training_quantization)

for further information on the types available and how to use them.

OPTIMIZE_FOR_SIZE

Deprecated. Does the same as DEFAULT.

OPTIMIZE_FOR_LATENCY

Deprecated. Does the same as DEFAULT.

EXPERIMENTAL_SPARSITY

Experimental flag, subject to change.

Enable optimization by taking advantage of the sparse model weights

trained with pruning.

The converter will inspect the sparsity pattern of the model weights and

do its best to improve size and latency.

The flag can be used alone to optimize float32 models with sparse weights.

It can also be used together with the DEFAULT optimization mode to

optimize quantized models with sparse weights.

"""

# Default optimization strategy that quantizes model weights. Enhanced

# optimizations are gained by providing a representative dataset that

# quantizes biases and activations as well.

# Converter will do its best to reduce size and latency, while minimizing

# the loss in accuracy.

DEFAULT = "DEFAULT"

# Deprecated. Does the same as DEFAULT.

OPTIMIZE_FOR_SIZE = "OPTIMIZE_FOR_SIZE"

# Deprecated. Does the same as DEFAULT.

OPTIMIZE_FOR_LATENCY = "OPTIMIZE_FOR_LATENCY"

# Experimental flag, subject to change.

# Enable optimization by taking advantage of the sparse model weights trained

# with pruning.

#

# The converter will inspect the sparsity pattern of the model weights and do

# its best to improve size and latency.

# The flag can be used alone to optimize float32 models with sparse weights.

# It can also be used together with the DEFAULT optimization mode to optimize

# quantized models with sparse weights.

EXPERIMENTAL_SPARSITY = "EXPERIMENTAL_SPARSITY"

def __str__(self):

return str(self.value)

# TODO(b/198099651): move converter implementation out of lite.py

@_tf_export("lite.RepresentativeDataset")

class RepresentativeDataset:

"""Representative dataset used to optimize the model.

This is a generator function that provides a small dataset to calibrate or

estimate the range, i.e, (min, max) of all floating-point arrays in the model

(such as model input, activation outputs of intermediate layers, and model

output) for quantization. Usually, this is a small subset of a few hundred

samples randomly chosen, in no particular order, from the training or

evaluation dataset.

"""

def __init__(self, input_gen):

"""Creates a representative dataset.

Args:

input_gen: A generator function that generates input samples for the model

and has the same order, type and shape as the inputs to the model.

Usually, this is a small subset of a few hundred samples randomly

chosen, in no particular order, from the training or evaluation dataset.

"""

self.input_gen = input_gen

@_tf_export("lite.TargetSpec")

class TargetSpec:

"""Specification of target device used to optimize the model.

Attributes:

supported_ops: Experimental flag, subject to change. Set of `tf.lite.OpsSet`

options, where each option represents a set of operators supported by the

target device. (default {tf.lite.OpsSet.TFLITE_BUILTINS}))

supported_types: Set of `tf.dtypes.DType` data types supported on the target

device. If initialized, optimization might be driven by the smallest type

in this set. (default set())

experimental_select_user_tf_ops: Experimental flag, subject to change. Set

of user's TensorFlow operators' names that are required in the TensorFlow

Lite runtime. These ops will be exported as select TensorFlow ops in the

model (in conjunction with the tf.lite.OpsSet.SELECT_TF_OPS flag). This is

an advanced feature that should only be used if the client is using TF ops

that may not be linked in by default with the TF ops that are provided

when using the SELECT_TF_OPS path. The client is responsible for linking

these ops into the target runtime.

experimental_supported_backends: Experimental flag, subject to change. Set

containing names of supported backends. Currently only "GPU" is supported,

more options will be available later.

"""

def __init__(

self,

supported_ops=None,

supported_types=None,

experimental_select_user_tf_ops=None,

experimental_supported_backends=None,

):

if supported_ops is None:

supported_ops = {OpsSet.TFLITE_BUILTINS}

self.supported_ops = supported_ops

if supported_types is None:

supported_types = set()

self.supported_types = supported_types

if experimental_select_user_tf_ops is None:

experimental_select_user_tf_ops = set()

self.experimental_select_user_tf_ops = experimental_select_user_tf_ops

self.experimental_supported_backends = experimental_supported_backends

self._experimental_custom_op_registerers = []

# Hint for the supported accumulation type used for inference. Typically

# used for fp16 post-training quantization, where some models can use fp16

# accumulators instead of the typical fp32 type.

self._experimental_supported_accumulation_type = None

class QuantizationMode:

"""QuantizationMode determines the quantization type from user options."""

def __init__(

self,

optimizations,

target_spec,

representative_dataset,

graph_def,

disable_per_channel=False,

experimental_new_dynamic_range_quantizer=False,

experimental_low_bit_qat=False,

full_integer_quantization_bias_type=None,

experimental_mlir_variable_quantization=False,

experimental_qdq_annotation=False,

):

self._optimizations = optimizations

for deprecated_optimization in [

Optimize.OPTIMIZE_FOR_SIZE,

Optimize.OPTIMIZE_FOR_LATENCY,

]:

if deprecated_optimization in self._optimizations:

logging.warning(

(

"Optimization option %s is deprecated, please use"

" optimizations=[Optimize.DEFAULT] instead."

),

deprecated_optimization,

)

self._experimental_qdq_annotation = experimental_qdq_annotation

self._target_spec = target_spec

self._representative_dataset = representative_dataset

self._graph_def = graph_def

if self._is_int8_target_required():

self._validate_int8_required()

self.enable_mlir_variable_quantization = (

experimental_mlir_variable_quantization

)

if self._is_float16_target_required():

self._validate_float16_required()

self._disable_per_channel = disable_per_channel

self._enable_new_dynamic_range_quantizer = (

experimental_new_dynamic_range_quantizer

)

# Allow training with lower than 8 bit weights to be converted

# to constants with trained scale.

self._experimental_low_bit_qat = experimental_low_bit_qat

self._full_integer_quantization_bias_type = (

full_integer_quantization_bias_type

)

self._validate_full_integer_quantization_bias_type()

def is_post_training_int8_only_quantization(self):

return (

self.is_any_optimization_enabled()

and self._representative_dataset is not None

and not self._is_int16x8_target_required()

and not self.is_allow_float()

and self._is_int8_target_required()

)

def is_post_training_int8_quantization_with_float_fallback(self):

return (

self.is_any_optimization_enabled()

and self._representative_dataset is not None

and not self._is_int16x8_target_required()

and self.is_allow_float()

and self._smallest_supported_type() == _dtypes.int8

)

def is_post_training_int8_quantization(self):

return (

self.is_post_training_int8_only_quantization()

or self.is_post_training_int8_quantization_with_float_fallback()

)

def is_post_training_int16x8_only_quantization(self):

return (

self.is_any_optimization_enabled()

and self._representative_dataset is not None

and self._is_int16x8_target_required()

and not self.is_allow_float()

)

def is_post_training_int16x8_quantization_with_float_fallback(self):

return (

self.is_any_optimization_enabled()

and self._representative_dataset is not None

and self._is_int16x8_target_required()

and self.is_allow_float()

)

def is_post_training_int16x8_quantization(self):

return (

self.is_post_training_int16x8_only_quantization()

or self.is_post_training_int16x8_quantization_with_float_fallback()

)

def is_post_training_integer_quantization(self):

return (

self.is_post_training_int8_quantization()

or self.is_post_training_int16x8_quantization()

)

def is_low_bit_quantize_aware_training(self):

return (

self.is_any_optimization_enabled()

and self.is_quantization_aware_trained_model()

and self._experimental_low_bit_qat

)

def is_quantization_aware_training(self):

if self._experimental_qdq_annotation:

return True

return (

self.is_any_optimization_enabled()

and self.is_quantization_aware_trained_model()

and not self.is_low_bit_quantize_aware_training()

)

def is_integer_quantization(self):

return (

self.is_post_training_integer_quantization()

or self.is_quantization_aware_training()

or self.is_low_bit_quantize_aware_training()

)

def is_post_training_dynamic_range_quantization(self):

# Post-training dynamic range quantization is only enabled if post-training

# int8 quantization and training time quantization was not done.

return (

self.is_any_optimization_enabled()

and self._representative_dataset is None

and not self.is_quantization_aware_trained_model()

and self._smallest_supported_type() == _dtypes.int8

)

def is_post_training_float16_quantization(self):

return (

self.is_any_optimization_enabled()

and self._smallest_supported_type().size == 2

and _dtypes.float16 in self._target_spec.supported_types

)

def is_bfloat16_quantization(self):

return (

self.is_any_optimization_enabled()

and self._smallest_supported_type().size == 2

and _dtypes.bfloat16 in self._target_spec.supported_types

)

def activations_type(self):

if self.is_integer_quantization():

if self._is_int16x8_target_required():

return _dtypes.int16

else:

return _dtypes.int8

else:

return _dtypes.float32

def bias_type(self):

if self._full_integer_quantization_bias_type:

return self._full_integer_quantization_bias_type

if self.activations_type() == _dtypes.int16:

return _dtypes.int64

elif self.activations_type() == _dtypes.int8:

return _dtypes.int32

else:

return _dtypes.float32

def converter_flags(self, inference_ty=None, inference_input_ty=None):

"""Flags to the converter."""

if self.is_integer_quantization():

is_low_bit_qat = self.is_low_bit_quantize_aware_training()

return {

"inference_type": (

inference_ty

if inference_ty is not None

else self.activations_type()

),

"inference_input_type": _dtypes.float32,

"post_training_quantize": False, # disable dynamic range quantization

"quantize_to_float16": False, # disable float16 quantization

"disable_infer_tensor_range": is_low_bit_qat,

"use_fake_quant_num_bits": is_low_bit_qat,

"enable_mlir_variable_quantization": (

self.enable_mlir_variable_quantization

),

}

elif self.is_post_training_dynamic_range_quantization():

return {

"inference_type": _dtypes.float32,

"inference_input_type": _dtypes.float32,

"post_training_quantize": True, # enable dynamic range quantization

"quantize_to_float16": False, # disable float16 quantization

# experimental: disable per-channel (per-axis) quantization.

"disable_per_channel_quantization": self._disable_per_channel,

"enable_mlir_dynamic_range_quantizer": (

self._enable_new_dynamic_range_quantizer

),

"enable_mlir_variable_quantization": (

self.enable_mlir_variable_quantization

),

}

elif self.is_post_training_float16_quantization():

return {

"inference_type": _dtypes.float32,

"inference_input_type": _dtypes.float32,

"post_training_quantize": True,

"quantize_to_float16": True, # enable float16 quantization

# pylint: disable=protected-access

"accumulation_type": (

self._target_spec._experimental_supported_accumulation_type

),

# pylint: enable=protected-access

"allow_bfloat16": self.is_bfloat16_quantization(),

"enable_mlir_dynamic_range_quantizer": (

self._enable_new_dynamic_range_quantizer

),

"enable_mlir_variable_quantization": (

self.enable_mlir_variable_quantization

),

}

else:

# Note this might still trigger (uint8) quantization to be compatible with

# the old converter.

return {

"inference_type": (

inference_ty if inference_ty is not None else _dtypes.float32

),

"inference_input_type": inference_input_ty,

"post_training_quantize": False, # enable dynamic range quantization

"quantize_to_float16": False, # disable float16 quantization

"allow_bfloat16": self.is_bfloat16_quantization(),

}

# Below are helpers for the above functions.

def _validate_int8_required(self):

"""Int8 mode requires certain parameters to exist and be compatible."""

# Validate target_spec attibute.

if set(self._target_spec.supported_ops) == {

OpsSet.TFLITE_BUILTINS_INT8

} and not (

set(self._target_spec.supported_types) == set()

or set(self._target_spec.supported_types) == {_dtypes.int8}

):

raise ValueError(

"As full integer quantization has been enabled by setting "

"`target_spec.supported_ops`={tf.lite.OpsSet.TFLITE_BUILTINS_INT8}, "

"thus `target_spec.supported_types` should be left uninitizalized "

"or set to {tf.int8}."

)

if set(self._target_spec.supported_types) == {_dtypes.int8}:

self._target_spec.supported_ops = {OpsSet.TFLITE_BUILTINS_INT8}

# Check if representative_dataset is specified.

if (

not self._representative_dataset

and not self.is_quantization_aware_training()

):

raise ValueError(

"For full integer quantization, a "

"`representative_dataset` must be specified."

)

# Update represenative dataset to the expected format.

if self._representative_dataset:

if not isinstance(self._representative_dataset, RepresentativeDataset):

self._representative_dataset = RepresentativeDataset(

self._representative_dataset

)

def _validate_float16_required(self):

"""Float16 mode requires certain parameters to exist and be compatible."""

if self.enable_mlir_variable_quantization:

raise ValueError(

"`_experimental_variable_quantization` is only supported for full"

" integer quantization."

)

def _validate_full_integer_quantization_bias_type(self):

"""Validates bias type for full interger quantization."""

bias_type = self._full_integer_quantization_bias_type

if not bias_type:

return

if self.activations_type() == _dtypes.float32:

raise ValueError(

"`full_integer_quantization_bias_type` is only supported for full"

" integer quantization."

)

if self.activations_type() == _dtypes.int8 and bias_type != _dtypes.int32:

raise ValueError(

"Expected bias type to be `dtypes.int32` for Int8Quant. "

f"Current setting bias type: {bias_type}"

)

if (

self.activations_type() == _dtypes.int16

and bias_type != _dtypes.int32

and bias_type != _dtypes.int64

):

raise ValueError(

"Expected bias type to be `dtypes.int32` or `dtypes.int64` for "

f"Int16Quant. Current setting bias type: {bias_type}"

)

def _is_int8_target_required(self):

return (

OpsSet.TFLITE_BUILTINS_INT8 in set(self._target_spec.supported_ops)

) or (set(self._target_spec.supported_types) == set([_dtypes.int8]))

def _is_int16x8_target_required(self):

return (

OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8

in set(self._target_spec.supported_ops)

)

def is_allow_float(self):

return (OpsSet.TFLITE_BUILTINS in set(self._target_spec.supported_ops)) or (

OpsSet.SELECT_TF_OPS in set(self._target_spec.supported_ops)

)

def _is_float16_target_required(self):

return _dtypes.float16 in self._target_spec.supported_types

def is_any_optimization_enabled(self):

return bool(

set(self._optimizations).intersection([

Optimize.OPTIMIZE_FOR_LATENCY,

Optimize.OPTIMIZE_FOR_SIZE,

Optimize.DEFAULT,

])

)

def _smallest_supported_type(self):

if self._target_spec.supported_types:

return min(self._target_spec.supported_types, key=lambda x: x.size)

else:

# The default smallest supported type is INT8.

return _dtypes.int8

def is_quantization_aware_trained_model(self):

"""Checks if the graph contains any training-time quantization ops."""

training_quant_ops = frozenset({

"FakeQuantWithMinMaxVars",

"FakeQuantWithMinMaxVarsPerChannel",

"FakeQuantWithMinMaxArgs",

"QuantizeAndDequantizeV2",

"QuantizeAndDequantizeV3",

})

if self._graph_def:

for node_def in self._graph_def.node:

if node_def.op in training_quant_ops:

return True

for function in self._graph_def.library.function:

for node_def in function.node_def:

if node_def.op in training_quant_ops:

return True

return False

class TFLiteConverterBase:

"""Converter superclass to share functionality between V1 and V2 converters."""

# Stores the original model type temporarily to transmit the information

# from the factory class methods to TFLiteConverterBase init function.

_original_model_type = conversion_metadata_fb.ModelType.NONE

def __init__(self):

self.optimizations = set()

self.representative_dataset = None

self.target_spec = TargetSpec()

self.allow_custom_ops = False

self.experimental_new_converter = True

self.experimental_new_quantizer = True

self.experimental_enable_resource_variables = True

self._experimental_calibrate_only = False

self._experimental_sparsify_model = False

self._experimental_disable_per_channel = False

self._debug_info = None # contains the stack traces of all the original

# nodes in the `GraphDef` to the converter.

self.saved_model_dir = None

self._saved_model_tags = None

self._saved_model_version = 0

self._saved_model_exported_names = []

self._tflite_metrics = metrics.TFLiteConverterMetrics()

self._collected_converter_params = {}

self.unfold_batchmatmul = False

self.legalize_custom_tensor_list_ops = False

self._experimental_lower_tensor_list_ops = True

self._experimental_default_to_single_batch_in_tensor_list_ops = False

self._experimental_unfold_large_splat_constant = False

self._experimental_tf_quantization_mode = None

# If unset, bias:int32 is by default except 16x8 quant.

# For 16x8 quant, bias:int64 is used to prevent any overflow by default.

# The accumulator type will be the same as bias type set by

# full_integer_quantization_bias_type.

self._experimental_full_integer_quantization_bias_type = None

# Provides specs for quantization, whether preset or custom.

self._experimental_quantization_options = None # Deprecated

# Whether to use StableHLO Quantizer instead of TFLite Quantizer.

self.experimental_use_stablehlo_quantizer = False

# Quantization configuration to pass to StableHLO Quantizer.

self.experimental_stablehlo_quantizer_config = None

# Initializes conversion metadata.

self.exclude_conversion_metadata = False

self._metadata = conversion_metadata_fb.ConversionMetadataT()

self._metadata.environment = conversion_metadata_fb.EnvironmentT()

self._metadata.options = conversion_metadata_fb.ConversionOptionsT()

self._metadata.environment.tensorflowVersion = versions.__version__

self._metadata.environment.modelType = self._get_original_model_type()

self._experimental_enable_dynamic_update_slice = False

self._experimental_preserve_assert_op = False

self._experimental_guarantee_all_funcs_one_use = False

# When the value is true, the MLIR quantantizer triggers dynamic range

# quantization in MLIR instead of the old quantizer. Used only if

# experimental_new_quantizer is on.

self.experimental_new_dynamic_range_quantizer = True

# Experimental flag to enable low-bit QAT in 8 bit.

self._experimental_low_bit_qat = False

# Experimental flag to add all TF ops (including custom TF ops) to the

# converted model as flex ops.

self._experimental_allow_all_select_tf_ops = False

self._experimental_variable_quantization = False

self._experimental_disable_fuse_mul_and_fc = False

self._experimental_use_buffer_offset = False

self._experimental_reduce_type_precision = False

self._experimental_qdq_conversion_mode = None

self._experimental_disable_per_channel_quantization_for_dense_layers = False

self._experimental_enable_composite_direct_lowering = True

self.model_origin_framework = constants.UNSET

self.canonicalizing_inf_as_min_max_float = True

self._experimental_strict_qdq = False

self._experimental_unsafe_fuse_dynamic_shaped_broadcast = False

self._experimental_unsafe_single_batch_rank_reduction = False

# Debug parameters

self.ir_dump_dir = None

self.ir_dump_pass_regex = None

self.ir_dump_func_regex = None

self.enable_timing = None

self.print_ir_before = None

self.print_ir_after = None

self.print_ir_module_scope = None

self.elide_elementsattrs_if_larger = None

self.serialize_debug_metadata = False

def _grappler_config(self, optimizers=None):

"""Creates a tf.compat.v1.ConfigProto for configuring Grappler.

Args:

optimizers: List of strings that represents the list of optimizers.

Returns:

tf.ConfigProto.

"""

if not optimizers:

optimizers = []

# MLIR converter will take care of constant folding instead of grappler.

if not self.experimental_new_converter:

optimizers.append("constfold")

is_only_flex_enabled = set([OpsSet.SELECT_TF_OPS]) == set(

self.target_spec.supported_ops

)

if is_only_flex_enabled:

# The layout optimizer turns NHCW to NCHW. This provides performance

# optimizations when Flex mode is enabled. However, this is not compatible

# with builtin ops.

optimizers.append("layout")

return _get_grappler_config(optimizers)

def _quantize(

self,

result,

input_type,

output_type,

activations_type,

bias_type,

allow_float,

enable_variable_quantization,

debug_options,

):

"""Quantize the model."""

# pylint: disable=protected-access

custom_op_registerers_by_name = [

x

for x in self.target_spec._experimental_custom_op_registerers

if isinstance(x, str)

]

custom_op_registerers_by_func = [

x

for x in self.target_spec._experimental_custom_op_registerers

if not isinstance(x, str)

]

# pylint: enable=protected-access

if not isinstance(self.representative_dataset, RepresentativeDataset):

self.representative_dataset = RepresentativeDataset(

self.representative_dataset

)

# Add intermediate tensors to the model if needed.

result = _calibrator.add_intermediate_tensors(result)

calibrate_quantize = _calibrator.Calibrator(

result, custom_op_registerers_by_name, custom_op_registerers_by_func

)

if self._experimental_calibrate_only or self.experimental_new_quantizer:

calibrated = calibrate_quantize.calibrate(

self.representative_dataset.input_gen

)

if self._experimental_calibrate_only:

return calibrated

elif self.experimental_new_quantizer and (

activations_type != _dtypes.int16

):

return _mlir_quantize(

calibrated,

self._experimental_disable_per_channel,

input_data_type=input_type,

output_data_type=output_type,

enable_variable_quantization=enable_variable_quantization,

disable_per_channel_for_dense_layers=self._experimental_disable_per_channel_quantization_for_dense_layers,

debug_options_str=debug_options.SerializeToString(),

)

else:

return calibrate_quantize.calibrate_and_quantize(

self.representative_dataset.input_gen,

input_type,

output_type,

allow_float,

activations_type,

bias_type,

disable_per_channel=self._experimental_disable_per_channel,

disable_per_channel_quantization_for_dense_layers=self._experimental_disable_per_channel_quantization_for_dense_layers,

)

def _is_unknown_shapes_allowed(self):

# Unknown dimensions are only allowed with the new converter.

return self.experimental_new_converter

def _get_base_converter_args(self):

"""Returns the base converter args.

Returns:

{key str: val}

"""

args = {

"input_format": constants.TENSORFLOW_GRAPHDEF,

"allow_custom_ops": self.allow_custom_ops,

"debug_info": self._debug_info,

"target_ops": self.target_spec.supported_ops,

"select_user_tf_ops": self.target_spec.experimental_select_user_tf_ops,

"supported_backends": self.target_spec.experimental_supported_backends,

"unfold_batchmatmul": self.unfold_batchmatmul,

"legalize_custom_tensor_list_ops": self.legalize_custom_tensor_list_ops,

"lower_tensor_list_ops": self._experimental_lower_tensor_list_ops,

"unfold_large_splat_constant": (

self._experimental_unfold_large_splat_constant

),

"default_to_single_batch_in_tensor_list_ops": (

self._experimental_default_to_single_batch_in_tensor_list_ops

),

"tf_quantization_mode": self._experimental_tf_quantization_mode,

"experimental_enable_resource_variables": (

self.experimental_enable_resource_variables

),

"enable_dynamic_update_slice": (

self._experimental_enable_dynamic_update_slice

),

"preserve_assert_op": self._experimental_preserve_assert_op,

"guarantee_all_funcs_one_use": (

self._experimental_guarantee_all_funcs_one_use

),

"allow_all_select_tf_ops": self._experimental_allow_all_select_tf_ops,

"disable_fuse_mul_and_fc": self._experimental_disable_fuse_mul_and_fc,

"quantization_options": self._experimental_quantization_options,

"ir_dump_dir": self.ir_dump_dir,

"ir_dump_pass_regex": self.ir_dump_pass_regex,

"ir_dump_func_regex": self.ir_dump_func_regex,

"enable_timing": self.enable_timing,

"print_ir_before": self.print_ir_before,

"print_ir_after": self.print_ir_after,

"print_ir_module_scope": self.print_ir_module_scope,

"elide_elementsattrs_if_larger": self.elide_elementsattrs_if_larger,

"use_buffer_offset": self._experimental_use_buffer_offset,

"reduce_type_precision": self._experimental_reduce_type_precision,

"use_stablehlo_quantizer": self.experimental_use_stablehlo_quantizer,

"stablehlo_quantizer_config": (

self.experimental_stablehlo_quantizer_config

),

"qdq_conversion_mode": self._experimental_qdq_conversion_mode,

"strict_qdq_mode": self._experimental_strict_qdq,

"disable_per_channel_quantization_for_dense_layers": (

self._experimental_disable_per_channel_quantization_for_dense_layers

),

"enable_composite_direct_lowering": (

self._experimental_enable_composite_direct_lowering

),

"model_origin_framework": self.model_origin_framework,

"canonicalizing_inf_as_min_max_float": (

self.canonicalizing_inf_as_min_max_float

),

"serialize_debug_metadata": self.serialize_debug_metadata,

"unsafe_fuse_dynamic_shaped_broadcast": (

self._experimental_unsafe_fuse_dynamic_shaped_broadcast

),

"unsafe_single_batch_rank_reduction": (

self._experimental_unsafe_single_batch_rank_reduction

),

}

if self.saved_model_dir:

args.update({

"saved_model_dir": self.saved_model_dir,

"saved_model_version": self._saved_model_version,

"saved_model_tags": self._saved_model_tags,

"saved_model_exported_names": self._saved_model_exported_names,

})

if self._experimental_quantization_options:

logging.warning(

"Configs from custom methods in experimental_quantization_options"

" may not produce a valid tflite model. Note that currently this"

" option only supports StableHLO path. Setting this option in TFLite"

" path will be a no-op."

)

if self.experimental_use_stablehlo_quantizer:

self._assign_stablehlo_quantization_config_or_populate_default(args)

elif self.experimental_stablehlo_quantizer_config is not None:

raise ValueError(

"QuantizationConfig should be provided only when"

" experimental_use_stablehlo_quantizer is set to true."

)

return args

def _assign_stablehlo_quantization_config_or_populate_default(self, args):

"""Assigns `QuantizationConfig` to `args` or populate default.

Args:

args: Dictionary of argument names and associated values.

"""

if (

self.experimental_stablehlo_quantizer_config is not None

and Optimize.DEFAULT not in self.optimizations

):

args["quantization_config"] = self.experimental_stablehlo_quantizer_config

elif Optimize.DEFAULT in self.optimizations and self.representative_dataset:

if len(self._saved_model_exported_names) != 1:

raise ValueError(

"StableHLO quantizer is only supported when converting from a"

" SavedModel with one signature key."

)

signature_key = self._saved_model_exported_names[0]

# Convert a programmatically provided representative dataset to a

# temporary TFRecord file to be used by the StableHLO quantizer.

tfrecord_file_path = tempfile.mkstemp(

suffix=".tfrecord", prefix=signature_key

)[1]

rd.TfRecordRepresentativeDatasetSaver(

{signature_key: tfrecord_file_path}

).save({signature_key: self.representative_dataset()})

quantization_config = qc.QuantizationConfig(

static_range_ptq_preset=qc.StaticRangePtqPreset(

representative_datasets=[

qc.RepresentativeDatasetConfig(

tf_record=qc.TfRecordFile(path=tfrecord_file_path)

)

],

enable_per_channel_quantized_weight=True,

enable_full_int_quantization=True,

),

# For ODML use cases, uniform quantized types should be left intact.

pipeline_config=qc.PipelineConfig(

unpack_quantized_types=False,

),

)

args["quantization_config"] = quantization_config

# TODO: b/307626463 - Enable StableHLO quantizer DRQ when Optimize.DEFAULT

# is set without representative dataset.

else:

raise ValueError(

"StableHLO quantizer only supports static-range and weight-only PTQ."

)

def _contains_function_with_implements_attr(self, saved_model_proto):

meta_graph = saved_model_proto.meta_graphs[0]

for function in meta_graph.graph_def.library.function:

if function.attr.get("_implements", None) or function.attr.get(

"api_implements", None

):

return True

return False

def _parse_saved_model_args(self, always_enable_saved_model_import=False):

"""Parses SavedModel arguments from the given Keras/RNN SavedModel.

Args:

always_enable_saved_model_import: Bool. When the value is true, it enables

MLIR saved model import path regardless of checking the conditions.

"""

if not self.experimental_new_converter:

self.saved_model_dir = None

return

if self.saved_model_dir:

try:

saved_model_proto, _ = _parse_saved_model_with_debug_info(

self.saved_model_dir

)

except OSError:

# If it fails to read the given saved model, it will fall back to the

# frozen graph def path.

self.saved_model_dir = None

return

if (

not always_enable_saved_model_import

and not self._contains_function_with_implements_attr(

saved_model_proto

)

):

self.saved_model_dir = None

return

if not self._saved_model_exported_names:

self._saved_model_exported_names = []

self._saved_model_version = saved_model_proto.saved_model_schema_version

if self._saved_model_version == 0:

self.saved_model_dir = None

logging.warning("SavedModel schema version is zero.")

return

if self._saved_model_version not in [1, 2]:

raise ValueError(

"SavedModel file format({0}) is not supported".format(

self._saved_model_version

)

)

def _sparsify_model(self):

return Optimize.EXPERIMENTAL_SPARSITY in self.optimizations

def _increase_conversion_attempt_metric(self):