This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 0bc3a6fa5f [Relax][Frontend][TFLite] Add missing TFLite operator
mappings (#19813)
0bc3a6fa5f is described below
commit 0bc3a6fa5fcc6e73658f027f86e2385efd443198
Author: Hongyi Wu <[email protected]>
AuthorDate: Thu Jun 18 13:07:32 2026 +0800
[Relax][Frontend][TFLite] Add missing TFLite operator mappings (#19813)
## Summary
Adds Relax TFLite frontend coverage for several builtin operators whose
lowering can be expressed with existing Relax ops. The change is limited
to
the TFLite frontend converter and its tests.
## Changes
- Add direct TFLite operator mappings:
- `SIGN` -> `relax.op.sign`
- `BITWISE_XOR` -> `relax.op.bitwise_xor`
- `RIGHT_SHIFT` -> `relax.op.right_shift`
- `RELU_0_TO_1` -> `relax.op.clip`
- `BUCKETIZE` -> `relax.op.bucketize`
- Add `RANK` lowering as a scalar rank constant for statically-ranked
inputs.
- Add `UNIQUE` lowering through `relax.op.unique`, returning unique
values and
inverse indices.
- Extend existing scatter-based segment lowering for:
- `UNSORTED_SEGMENT_SUM`
- `UNSORTED_SEGMENT_MAX`
- Fix the `FAKE_QUANT` narrow-range vector path to build Relax constants
with
`relax.const`.
## Tests
Adds coverage in `tests/python/relax/test_frontend_tflite.py` for the
new
mappings and the `FAKE_QUANT` narrow-range vector path. Most tests use
structural-equal Relax IR checks; data-dependent and numeric regression
cases
add focused output checks for `UNIQUE` and `FAKE_QUANT`.
Local validation:
```bash
python -m ruff format \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest \
tests/python/relax/test_frontend_tflite.py \
-k "unique or sign or bitwise_xor or right_shift or bucketize or
relu_0_to_1 or rank or unsorted_segment_sum or unsorted_segment_max or range or
fake_quant or dynamic_update_slice"
python -m pytest tests/python/relax/test_frontend_tflite.py
```
Result:
- ruff format: 2 files left unchanged
- ruff check: All checks passed
- target pytest: 22 passed, 523 deselected
- full TFLite pytest: 545 passed
---
.../tvm/relax/frontend/tflite/tflite_frontend.py | 123 +++++++-
tests/python/relax/test_frontend_tflite.py | 318 +++++++++++++++++++++
2 files changed, 436 insertions(+), 5 deletions(-)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 0edfc00ce9..35e75e89bc 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -203,8 +203,10 @@ class OperatorConverter:
"BIDIRECTIONAL_SEQUENCE_LSTM":
self.convert_bidirectional_sequence_lstm,
"BIDIRECTIONAL_SEQUENCE_RNN":
self.convert_bidirectional_sequence_rnn,
"BITCAST": self.convert_bitcast,
- "BROADCAST_TO": self.convert_broadcast_to,
+ "BITWISE_XOR": functools.partial(self._convert_elemwise,
relax_op=_op.bitwise_xor),
"BROADCAST_ARGS": self.convert_broadcast_args,
+ "BROADCAST_TO": self.convert_broadcast_to,
+ "BUCKETIZE": self.convert_bucketize,
"CALL": self.convert_call,
"CALL_ONCE": self.convert_call_once,
"COMPLEX_ABS": self.convert_complex_abs,
@@ -293,6 +295,7 @@ class OperatorConverter:
"POW": functools.partial(self._convert_elemwise,
relax_op=_op.power),
"PRELU": self.convert_prelu,
"RANGE": self.convert_range,
+ "RANK": self.convert_rank,
"QUANTIZE": self.convert_quantize,
"RANDOM_STANDARD_NORMAL": self.convert_random_standard_normal,
"RANDOM_UNIFORM": self.convert_random_uniform,
@@ -304,12 +307,14 @@ class OperatorConverter:
"REDUCE_MIN": functools.partial(self._convert_reduce,
relax_op=_op.min),
"REDUCE_PROD": functools.partial(self._convert_reduce,
relax_op=_op.prod),
"RELU": self.convert_relu,
+ "RELU_0_TO_1": self.convert_relu_0_to_1,
"RELU6": self.convert_relu6,
"RELU_N1_TO_1": self.convert_relu_n1_to_1,
"RESHAPE": self.convert_reshape,
"RESIZE_BILINEAR": self.convert_resize_bilinear,
"RESIZE_NEAREST_NEIGHBOR": self.convert_resize_nearest_neighbor,
"RFFT2D": self.convert_rfft2d,
+ "RIGHT_SHIFT": functools.partial(self._convert_elemwise,
relax_op=_op.right_shift),
"ROUND": functools.partial(self._convert_unary_elemwise,
relax_op=_op.round),
"RSQRT": functools.partial(self._convert_unary_elemwise,
relax_op=_op.rsqrt),
"REVERSE_SEQUENCE": self.convert_reverse_sequence,
@@ -321,6 +326,7 @@ class OperatorConverter:
self._convert_segment_op, op_name="SEGMENT_SUM",
reduction="add"
),
"SHAPE": self.convert_shape,
+ "SIGN": functools.partial(self._convert_unary_elemwise,
relax_op=_op.sign),
"SIN": functools.partial(self._convert_unary_elemwise,
relax_op=_op.sin),
"SLICE": self.convert_slice,
"SOFTMAX": self.convert_softmax,
@@ -409,12 +415,19 @@ class OperatorConverter:
"TRANSPOSE": self.convert_transpose,
"UNPACK": self.convert_unpack,
"UNIDIRECTIONAL_SEQUENCE_RNN":
self.convert_unidirectional_sequence_rnn,
+ "UNIQUE": self.convert_unique,
+ "UNSORTED_SEGMENT_MAX": functools.partial(
+ self._convert_segment_op, op_name="UNSORTED_SEGMENT_MAX",
reduction="max"
+ ),
"UNSORTED_SEGMENT_MIN": functools.partial(
self._convert_segment_op, op_name="UNSORTED_SEGMENT_MIN",
reduction="min"
),
"UNSORTED_SEGMENT_PROD": functools.partial(
self._convert_segment_op, op_name="UNSORTED_SEGMENT_PROD",
reduction="mul"
),
+ "UNSORTED_SEGMENT_SUM": functools.partial(
+ self._convert_segment_op, op_name="UNSORTED_SEGMENT_SUM",
reduction="add"
+ ),
"UNIDIRECTIONAL_SEQUENCE_LSTM":
self.convert_unidirectional_sequence_lstm,
"VAR_HANDLE": self.convert_var_handle,
"WHERE": self.convert_select,
@@ -1563,6 +1576,18 @@ class OperatorConverter:
return out
+ def convert_rank(self, op):
+ """Convert TFLite RANK."""
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 1, "output tensors length should be 1"
+ output_dtype =
self.get_tensor_type_str(output_tensors[0].tensor.Type())
+
+ rank = len(self.get_tensor_shape(input_tensors[0]))
+ return relax.const(rank, dtype=output_dtype)
+
def convert_shape(self, op):
"""Convert TFLite Shape"""
@@ -1585,6 +1610,29 @@ class OperatorConverter:
return out
+ def convert_bucketize(self, op):
+ """Convert TFLite BUCKETIZE."""
+ from tflite.BucketizeOptions import BucketizeOptions
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.TensorType import TensorType
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.BucketizeOptions
+ op_options = op.BuiltinOptions()
+ bucketize_options = BucketizeOptions()
+ bucketize_options.Init(op_options.Bytes, op_options.Pos)
+
+ boundaries = self.bb.normalize(
+ relax.const(bucketize_options.BoundariesAsNumpy(), dtype="float32")
+ )
+ out_tensor = self.get_output_tensors(op)[0]
+ out_int32 = out_tensor.tensor.Type() == TensorType.INT32
+ return relax.op.bucketize(
+ self.get_tensor_expr(input_tensors[0]), boundaries,
out_int32=out_int32, right=False
+ )
+
def convert_relu(self, op):
"""Convert TFLite ReLU"""
@@ -1606,6 +1654,26 @@ class OperatorConverter:
return out
+ def convert_relu_0_to_1(self, op):
+ """Convert TFLite RELU_0_TO_1."""
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = input_tensors[0]
+ in_expr = self.get_expr(input_tensor.tensor_idx)
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 1, "output tensors length should be 1"
+ output_tensor = output_tensors[0]
+
+ if input_tensor.qnn_params:
+ in_f32 = self.dequantize(in_expr, input_tensor)
+ out = relax.op.clip(in_f32, 0, 1)
+ out = self.quantize(out, output_tensor)
+ else:
+ out = relax.op.clip(in_expr, 0, 1)
+
+ return out
+
def convert_hard_swish(self, op):
"""Convert TFLite Hard swish"""
input_tensors = self.get_input_tensors(op)
@@ -4898,12 +4966,57 @@ class OperatorConverter:
data = relax.op.zeros(shape, updates_dtype)
return relax.op.scatter_nd(data, indices, updates, "update")
+ def convert_unique(self, op):
+ """Convert TFLite UNIQUE."""
+ from tflite.TensorType import TensorType
+ from tflite.UniqueOptions import UniqueOptions
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "UNIQUE should have 1 input tensor"
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 2, "UNIQUE should have 2 output tensors"
+
+ unique_options = UniqueOptions()
+ op_options = op.BuiltinOptions()
+ unique_options.Init(op_options.Bytes, op_options.Pos)
+
+ unique = self.bb.normalize(
+ relax.op.unique(
+ self.get_tensor_expr(input_tensors[0]),
+ sorted=False,
+ return_index=False,
+ return_inverse=True,
+ return_counts=False,
+ axis=None,
+ )
+ )
+ values = self.bb.emit(relax.TupleGetItem(unique, 0))
+ inverse_indices = self.bb.emit(relax.TupleGetItem(unique, 1))
+
+ idx_out_type = unique_options.IdxOutType()
+ if idx_out_type == TensorType.INT32:
+ inverse_indices = self.bb.emit(relax.op.astype(inverse_indices,
"int32"))
+
+ return relax.Tuple([values, inverse_indices])
+
def _get_segment_scatter_base(self, output_shape, output_dtype, reduction):
"""Create the identity base tensor for scatter-based segment
reductions."""
if reduction == "add":
return relax.op.zeros(output_shape, output_dtype)
if reduction == "mul":
return relax.op.full(output_shape, relax.const(1, output_dtype),
output_dtype)
+ if reduction == "max":
+ np_dtype = np.dtype(output_dtype)
+ if np.issubdtype(np_dtype, np.floating):
+ identity = np.finfo(np_dtype).min
+ elif np.issubdtype(np_dtype, np.integer):
+ identity = np.iinfo(np_dtype).min
+ else:
+ raise tvm.error.OpNotImplemented(
+ f"UNSORTED_SEGMENT_MAX does not support output dtype
{output_dtype}."
+ )
+ return relax.op.full(output_shape, relax.const(identity,
output_dtype), output_dtype)
if reduction == "min":
np_dtype = np.dtype(output_dtype)
if np.issubdtype(np_dtype, np.floating):
@@ -7689,13 +7802,13 @@ class OperatorConverter:
nudged_min = (quant_min - nudged_zero_point) * scale
nudged_max = (quant_max - nudged_zero_point) * scale
- nudged_min_expr = relax.op.const(nudged_min)
+ nudged_min_expr = relax.const(nudged_min, "float32")
clamped = relax.op.clip(in_expr, nudged_min, nudged_max)
clamped_shifted = relax.op.subtract(clamped, nudged_min_expr)
- half = relax.op.const(0.5)
- one = relax.op.const(1.0)
- scale_expr = relax.op.const(scale)
+ half = relax.const(0.5, "float32")
+ one = relax.const(1.0, "float32")
+ scale_expr = relax.const(scale, "float32")
inv_scale = relax.op.divide(one, scale_expr)
rounded = relax.op.floor(_op.add(_op.multiply(clamped_shifted,
inv_scale), half))
return relax.op.add(_op.multiply(rounded, scale_expr), nudged_min_expr)
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index 01beaadd7b..5a4460d0b7 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -443,6 +443,99 @@ def test_bitcast_int16_to_int32_collapses_shape():
verify(BitcastI16ToI32, Expected)
+def test_bitwise_xor():
+ """BITWISE_XOR lowers to relax.op.bitwise_xor."""
+
+ class BitwiseXor(tf.Module):
+ @tf.function(
+ input_signature=[
+ tf.TensorSpec(shape=(2, 3), dtype=tf.int32),
+ tf.TensorSpec(shape=(2, 3), dtype=tf.int32),
+ ]
+ )
+ def func(self, x, y):
+ return tf.bitwise.bitwise_xor(x, y)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), dtype="int32"),
+ y: R.Tensor((2, 3), dtype="int32"),
+ ) -> R.Tensor((2, 3), dtype="int32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv: R.Tensor((2, 3), dtype="int32") = R.bitwise_xor(x, y)
+ R.output(gv)
+ return gv
+
+ verify(BitwiseXor, Expected)
+
+
+def test_right_shift():
+ """RIGHT_SHIFT lowers to relax.op.right_shift."""
+
+ class RightShift(tf.Module):
+ @tf.function(
+ input_signature=[
+ tf.TensorSpec(shape=(2, 3), dtype=tf.int32),
+ tf.TensorSpec(shape=(2, 3), dtype=tf.int32),
+ ]
+ )
+ def func(self, x, y):
+ return tf.bitwise.right_shift(x, y)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), dtype="int32"),
+ y: R.Tensor((2, 3), dtype="int32"),
+ ) -> R.Tensor((2, 3), dtype="int32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv: R.Tensor((2, 3), dtype="int32") = R.right_shift(x, y)
+ R.output(gv)
+ return gv
+
+ verify(RightShift, Expected)
+
+
+def test_sign():
+ """SIGN lowers to relax.op.sign."""
+
+ class Sign(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3),
dtype=tf.float32)])
+ def func(self, x):
+ return tf.sign(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 3), dtype="float32") = R.sign(x)
+ R.output(gv)
+ return gv
+
+ verify(Sign, Expected)
+
+
+def test_unique():
+ """UNIQUE returns values and inverse indices."""
+
+ class Unique(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(6,),
dtype=tf.int32)])
+ def func(self, x):
+ return tf.raw_ops.Unique(x=x, out_idx=tf.int64)
+
+ mod = _get_mod_from_cfunc(Unique().func.get_concrete_function())
+ values, inverse_indices = _run_module(mod, np.array([3, 1, 3, 2, 1, 2],
dtype=np.int32))
+ np.testing.assert_array_equal(values, np.array([3, 1, 2], dtype=np.int32))
+ np.testing.assert_array_equal(inverse_indices, np.array([0, 1, 0, 2, 1,
2], dtype=np.int64))
+
+
def test_expand_dims():
class ExpandDims(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=(1, 30),
dtype=tf.float32)])
@@ -533,6 +626,74 @@ def test_shape_dynamic_dim():
verify(ShapeDynamic)
+def _build_rank_model():
+ """Build a minimal TFLite RANK model."""
+ builder = flatbuffers.Builder(1024)
+ builtin_op = _get_builtin_operator("RANK")
+ op_code = _build_operator_code(builder, builtin_op)
+ options = _build_empty_builtin_options(builder, "RankOptions")
+
+ tensors = [
+ _build_tensor(builder, 0, [2, 3, 4]),
+ _build_tensor(builder, 1, [], tensor_type=_tfl_tensor_type.INT32),
+ ]
+ op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1],
+ builtin_options_type=_get_builtin_options_type("RankOptions"),
+ builtin_options=options,
+ )
+ subgraph = _build_subgraph(builder, tensors=tensors, operators=[op],
inputs=[0], outputs=[1])
+ return _finish_tflite_model(
+ builder,
+ subgraph=subgraph,
+ operator_codes=[op_code],
+ buffers=[_build_buffer(builder), _build_buffer(builder)],
+ )
+
+
+def test_rank():
+ """RANK emits a static rank constant."""
+ mod = _load_model_from_buffer(_build_rank_model())
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), dtype="float32")) -> R.Tensor((),
dtype="int32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((), dtype="int32") = R.const(3, "int32")
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_bucketize():
+ """BUCKETIZE lowers to relax.op.bucketize."""
+
+ class Bucketize(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3),
dtype=tf.float32)])
+ def func(self, x):
+ return tf.raw_ops.Bucketize(input=x, boundaries=[0.0, 1.0, 3.0])
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3),
dtype="int32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 3), dtype="int32") = R.bucketize(
+ x, R.const([0.0, 1.0, 3.0], "float32"), out_int32=True,
right=False
+ )
+ R.output(gv)
+ return gv
+
+ verify(Bucketize, Expected)
+
+
@pytest.mark.parametrize(
"start, limit, delta, dtype",
[
@@ -2085,6 +2246,66 @@ def test_unsorted_segment_min():
verify(Model, Expected)
+def test_unsorted_segment_sum():
+ """UNSORTED_SEGMENT_SUM lowers to scatter_nd with add reduction."""
+
+ class Model(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(4, 2),
dtype=tf.float32)])
+ def func(self, data):
+ return tf.raw_ops.UnsortedSegmentSum(
+ data=data,
+ segment_ids=tf.constant([0, 2, 1, 2], dtype=tf.int32),
+ num_segments=tf.constant(3, dtype=tf.int32),
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(data: R.Tensor((4, 2), dtype="float32")) -> R.Tensor((3, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ lv: R.Tensor((3, 2), dtype="float32") = R.zeros(R.shape([3,
2]), dtype="float32")
+ lv1: R.Tensor((4, 1), dtype="int32") = R.expand_dims(
+ R.const([0, 2, 1, 2], "int32"), axis=[1]
+ )
+ gv: R.Tensor((3, 2), dtype="float32") = R.scatter_nd(lv, lv1,
data, reduction="add")
+ R.output(gv)
+ return gv
+
+ verify(Model, Expected)
+
+
+def test_unsorted_segment_max():
+ """UNSORTED_SEGMENT_MAX lowers to scatter_nd with max reduction."""
+
+ class Model(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(4, 2),
dtype=tf.float32)])
+ def func(self, data):
+ return tf.raw_ops.UnsortedSegmentMax(
+ data=data,
+ segment_ids=tf.constant([0, 2, 1, 2], dtype=tf.int32),
+ num_segments=tf.constant(3, dtype=tf.int32),
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(data: R.Tensor((4, 2), dtype="float32")) -> R.Tensor((3, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ lv: R.Tensor((3, 2), dtype="float32") = R.full(
+ R.shape([3, 2]), R.const(np.finfo(np.float32).min,
"float32"), dtype="float32"
+ )
+ lv1: R.Tensor((4, 1), dtype="int32") = R.expand_dims(
+ R.const([0, 2, 1, 2], "int32"), axis=[1]
+ )
+ gv: R.Tensor((3, 2), dtype="float32") = R.scatter_nd(lv, lv1,
data, reduction="max")
+ R.output(gv)
+ return gv
+
+ verify(Model, Expected)
+
+
def test_unsorted_segment_prod():
"""UNSORTED_SEGMENT_PROD lowers to scatter_nd with mul reduction."""
@@ -3452,6 +3673,42 @@ def test_hard_swish():
verify(HardSwish, Expected)
+def _build_relu_0_to_1_model():
+ """Build a minimal TFLite RELU_0_TO_1 model."""
+ builder = flatbuffers.Builder(1024)
+ builtin_op = _get_builtin_operator("RELU_0_TO_1")
+ op_code = _build_operator_code(builder, builtin_op)
+ tensors = [
+ _build_tensor(builder, 0, [2, 2]),
+ _build_tensor(builder, 1, [2, 2]),
+ ]
+ op = _build_operator(builder, 0, [0], [1])
+ subgraph = _build_subgraph(builder, tensors=tensors, operators=[op],
inputs=[0], outputs=[1])
+ return _finish_tflite_model(
+ builder,
+ subgraph=subgraph,
+ operator_codes=[op_code],
+ buffers=[_build_buffer(builder), _build_buffer(builder)],
+ )
+
+
+def test_relu_0_to_1():
+ """RELU_0_TO_1 lowers to clip(0, 1)."""
+ mod = _load_model_from_buffer(_build_relu_0_to_1_model())
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 2), dtype="float32") = R.clip(x, min=0, max=1)
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_relu_n1_to_1():
class ReLU_N1_to_1(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=(1, 30),
dtype=tf.float32)])
@@ -3471,6 +3728,67 @@ def test_relu_n1_to_1():
verify(ReLU_N1_to_1, Expected)
+def _build_fake_quant_model(*, narrow_range, num_bits=8, min_value=-1.0,
max_value=1.0):
+ """Build a minimal TFLite FAKE_QUANT model."""
+ fake_quant_options = _get_tflite_schema_module("FakeQuantOptions")
+ builder = flatbuffers.Builder(1024)
+ builtin_op = _get_builtin_operator("FAKE_QUANT")
+ op_code = _build_operator_code(builder, builtin_op)
+
+ fake_quant_options.FakeQuantOptionsStart(builder)
+ fake_quant_options.FakeQuantOptionsAddMin(builder, min_value)
+ fake_quant_options.FakeQuantOptionsAddMax(builder, max_value)
+ fake_quant_options.FakeQuantOptionsAddNumBits(builder, num_bits)
+ fake_quant_options.FakeQuantOptionsAddNarrowRange(builder, narrow_range)
+ options = fake_quant_options.FakeQuantOptionsEnd(builder)
+
+ tensors = [
+ _build_tensor(builder, 0, [4]),
+ _build_tensor(builder, 1, [4]),
+ ]
+ op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1],
+ builtin_options_type=_get_builtin_options_type("FakeQuantOptions"),
+ builtin_options=options,
+ )
+ subgraph = _build_subgraph(builder, tensors=tensors, operators=[op],
inputs=[0], outputs=[1])
+ return _finish_tflite_model(
+ builder,
+ subgraph=subgraph,
+ operator_codes=[op_code],
+ buffers=[_build_buffer(builder), _build_buffer(builder)],
+ )
+
+
+def _fake_quant_reference(data, *, narrow_range, num_bits=8, min_value=-1.0,
max_value=1.0):
+ quant_min = 1 if narrow_range else 0
+ quant_max = (1 << num_bits) - 1
+ scale = (max_value - min_value) / (quant_max - quant_min)
+ zero_point_from_min = quant_min - min_value / scale
+ if zero_point_from_min <= quant_min:
+ nudged_zero_point = quant_min
+ elif zero_point_from_min >= quant_max:
+ nudged_zero_point = quant_max
+ else:
+ nudged_zero_point = round(zero_point_from_min)
+ nudged_min = (quant_min - nudged_zero_point) * scale
+ nudged_max = (quant_max - nudged_zero_point) * scale
+ clamped = np.clip(data, nudged_min, nudged_max)
+ return np.floor((clamped - nudged_min) / scale + 0.5) * scale + nudged_min
+
+
+def test_fake_quant_narrow_range_vector():
+ """FAKE_QUANT supports narrow_range on vector inputs."""
+ mod = _load_model_from_buffer(_build_fake_quant_model(narrow_range=True))
+ data = np.array([-2.0, -0.5, 0.5, 2.0], dtype=np.float32)
+ output = _run_module(mod, data)
+ expected = _fake_quant_reference(data,
narrow_range=True).astype(np.float32)
+ np.testing.assert_allclose(output, expected, rtol=1e-6, atol=1e-6)
+
+
def test_prelu_basic():
alpha_init = tf.keras.initializers.Constant(np.linspace(0.1, 0.3, 30,
dtype=np.float32))
prelu = tf.keras.layers.PReLU(alpha_initializer=alpha_init)
