ai-edge-quantizer-nightly 0.3.0.dev20250622__py3-none-any.whl → 0.3.0.dev20250624__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -113,6 +113,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     ),
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
     _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
+    _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -252,6 +253,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     ),
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
     _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
+    _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -745,6 +745,23 @@ def materialize_resize_bilinear(
   )
 
 
+def materialize_gather_nd(
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize tensors in tfl.gather_nd."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+      constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+      inputs_to_ignore=[1],  # Gather indices do not need to be quantized.
+  )
+
+
 def _get_tensor_shape_for_blockwise(
     tensor_shape: Sequence[int], quantized_dim: int, block_size: int
 ) -> list[int]:

ai_edge_quantizer/algorithms/utils/common_utils.py

@@ -293,17 +293,46 @@ def _materialize_standard_op_with_same_as_input_scale(
       get_tensor_quant_params_fn=get_tensor_quant_params_fn,
   )
   op_tensor_params.append(input_tensor_params)
-  # Use input quantization params for all output tensors.
-  _materialize_op_tensors(
-      op_tensor_params,
-      output_tensors,
-      is_inbounding_tensor=False,
-      op_info=op_info,
-      graph_info=graph_info,
-      tensor_name_to_qsv=tensor_name_to_qsv,
-      get_tensor_quant_params_fn=get_tensor_quant_params_fn,
-      quant_params=input_tensor_params.consumers[0].parameters,
+  # Use input quantization params for all output tensors but without
+  # quantized_data in case the input is a constant tensor.
+  input_quant_params = dataclasses.replace(
+      input_tensor_params.consumers[0].parameters,
+      quantized_data=None,
   )
+  if not isinstance(input_quant_params, qtyping.UniformQuantParams):
+    raise ValueError(
+        "_materialize_standard_op_with_same_as_input_scale only supports"
+        f" UniformQuantParams. For tensor {input_tensor_params.tensor_name},"
+        f" got {type(input_quant_params)}"
+    )
+  # Materialize each of the output tensors separately in case there are
+  # constants among them, requiring updating `quantized_data` first.
+  for output_tensor in output_tensors:
+    output_tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+        output_tensor, graph_info.buffers
+    )
+    # Quantize constant inputs' data with the output quantization params.
+    if output_tensor_data is None:
+      quant_params = input_quant_params
+    else:
+      quantized_data = uniform_quantize_tensor.uniform_quantize(
+          output_tensor_data, input_quant_params
+      )
+      quant_params = dataclasses.replace(
+          input_quant_params,
+          quantized_data=quantized_data,
+      )
+    _materialize_op_tensors(
+        op_tensor_params,
+        [output_tensor],
+        is_inbounding_tensor=False,
+        op_info=op_info,
+        graph_info=graph_info,
+        tensor_name_to_qsv=tensor_name_to_qsv,
+        get_tensor_quant_params_fn=get_tensor_quant_params_fn,
+        quant_params=quant_params,
+    )
+
   # Change output qsv to be the same as input qsv. This is safe since TFL
   # subgraph is acyclic.
   input_tensor_qsv = tensor_name_to_qsv[input_tensor_params.tensor_name]
@@ -841,13 +870,6 @@ def get_tensor_transformations(
       transformations = [_QuantTransformation.QUANTIZE_TENSOR]
     else:
       transformations = [_QuantTransformation.NO_QUANTIZE]
-  elif (
-      op_quant_config.weight_tensor_config is not None
-      and op_quant_config.weight_tensor_config.granularity
-      == qtyping.QuantGranularity.BLOCKWISE
-      and is_constant
-  ):
-    transformations = [_QuantTransformation.EMULATED_SUBCHANNEL]
   # Check if WEIGHT_ONLY.
   elif (
       op_quant_config.compute_precision == qtyping.ComputePrecision.FLOAT

ai_edge_quantizer/default_policy.py

@@ -186,7 +186,8 @@ DEFAULT_JSON_POLICY = """
       "STABLEHLO_COMPOSITE",
       "PAD",
       "MAX_POOL_2D",
-      "RESIZE_BILINEAR"
+      "RESIZE_BILINEAR",
+      "GATHER_ND"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -221,7 +222,8 @@ DEFAULT_JSON_POLICY = """
       "PAD",
       "SQUARED_DIFFERENCE",
       "MAX_POOL_2D",
-      "RESIZE_BILINEAR"
+      "RESIZE_BILINEAR",
+      "GATHER_ND"
     ],
     "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
     "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],

ai_edge_quantizer/qtyping.py

@@ -66,6 +66,7 @@ class TFLOperationName(str, enum.Enum):
   SQUARED_DIFFERENCE = 'SQUARED_DIFFERENCE'
   MAX_POOL_2D = 'MAX_POOL_2D'
   RESIZE_BILINEAR = 'RESIZE_BILINEAR'
+  GATHER_ND = 'GATHER_ND'
 
 
 class QuantizeMode(enum.Enum):
@@ -110,8 +111,8 @@ class QuantTransformation(enum.Enum):
   ADD_DEQUANTIZE = 2
   # Quantize the float tensor: float_tensor -> quantized_tensor.
   QUANTIZE_TENSOR = 3
-  # Create pattern for emulated subchannel quantization, only support fully
-  # connected op.
+  # (Deprecated) Create pattern for emulated subchannel quantization,
+  # only support fully connected op.
   EMULATED_SUBCHANNEL = 4
   # Duplicate the buffer.
   DUPLICATE_BUFFER = 5

ai_edge_quantizer/transformation_instruction_generator.py

@@ -673,7 +673,6 @@ class TransformationInstructionsGenerator:
     """
     is_tensor_unquantized = False
     is_tensor_quantized = False
-    is_operator_emulated = False
     for instruction in instructions:
       transform_type = instruction.transformation
       if transform_type == qtyping.QuantTransformation.NO_QUANTIZE:
@@ -683,17 +682,10 @@ class TransformationInstructionsGenerator:
           or transform_type == qtyping.QuantTransformation.ADD_DEQUANTIZE
       ):
         is_tensor_quantized = True
-      elif transform_type == qtyping.QuantTransformation.EMULATED_SUBCHANNEL:
-        is_operator_emulated = True
     if is_tensor_unquantized and is_tensor_quantized:
       raise ValueError(
           "Tensor %s can not be both quantized and unquantized" % tensor_name
       )
-    if is_operator_emulated and len(instructions) > 1:
-      raise ValueError(
-          "Tensor %s : op replacement transformation can not be combined with"
-          " other transformations." % tensor_name
-      )
 
   def _check_tensor_transformation_instructions_valid(
       self,

ai_edge_quantizer/transformation_instruction_generator_test.py

@@ -953,33 +953,6 @@ class InstructionGeneratorTest(parameterized.TestCase):
         instructions["StatefulPartitionedCall:0"], output_transformation
     )
 
-  def test_raise_error_on_op_replacement_transformation_is_not_unique(self):
-    test_model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/insert_dequant_test.tflite"
-    )
-    quant_parameters = {}
-    quant_parameters["tfl.quantize"] = qtyping.TensorTransformationParams(
-        "tfl.quantize",
-        qtyping.OpToTensorParams(
-            subgraph_op_id=0,
-            transformations=[
-                qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                qtyping.QuantTransformation.EMULATED_SUBCHANNEL,
-            ],
-            parameters=qtyping.UniformQuantParams(
-                8, None, np.array([1]), np.array([0])
-            ),
-        ),
-        [],
-    )
-    ins_gen = instruction_generator.TransformationInstructionsGenerator(
-        test_model_path
-    )
-    with self.assertRaisesRegex(
-        ValueError, "op replacement transformation can not be combined"
-    ):
-      ins_gen.quant_params_to_transformation_insts(quant_parameters)
-
   def test_raise_error_on_no_quant_conflict(self):
     test_model_path = os.path.join(
         TEST_DATA_PREFIX_PATH, "tests/models/insert_dequant_test.tflite"

ai_edge_quantizer/transformation_performer.py

@@ -24,7 +24,6 @@ from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.transformations import dequant_insert
 from ai_edge_quantizer.transformations import duplicate_buffer
 from ai_edge_quantizer.transformations import duplicate_tensor
-from ai_edge_quantizer.transformations import emulated_subchannel
 from ai_edge_quantizer.transformations import insert_hadamard_rotation
 from ai_edge_quantizer.transformations import quant_insert
 from ai_edge_quantizer.transformations import quantize_tensor
@@ -72,7 +71,7 @@ class TransformationPerformer:
             quantize_tensor.quantize_tensor
         ),
         qtyping.QuantTransformation.EMULATED_SUBCHANNEL: (
-            emulated_subchannel.emulated_subchannel
+            transformation_utils.raise_deprecated_error
         ),
         qtyping.QuantTransformation.ADD_QUANTIZE: quant_insert.insert_quant,
         qtyping.QuantTransformation.DUPLICATE_BUFFER: (

ai_edge_quantizer/transformations/transformation_utils.py

@@ -203,3 +203,10 @@ def add_new_activation_tensor(
   new_tensor_id = len(subgraph.tensors)
   subgraph.tensors.append(new_tensor)
   return new_tensor_id
+
+
+def raise_deprecated_error(_: TransformationInput):
+  raise NotImplementedError(
+      'This transformation is deprecated. Please contact AI Edge Quantizer team'
+      ' if you see this error.'
+  )

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -60,6 +60,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.SQUARED_DIFFERENCE: schema.BuiltinOperator.SQUARED_DIFFERENCE,
     _TFLOpName.MAX_POOL_2D: schema.BuiltinOperator.MAX_POOL_2D,
     _TFLOpName.RESIZE_BILINEAR: schema.BuiltinOperator.RESIZE_BILINEAR,
+    _TFLOpName.GATHER_ND: schema.BuiltinOperator.GATHER_ND,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.3.0.dev20250622
+Version: 0.3.0.dev20250624
 Summary: A quantizer for advanced developers to quantize converted AI Edge models.
 Home-page: https://github.com/google-ai-edge/ai-edge-quantizer
 Keywords: On-Device ML,AI,Google,TFLite,Quantization,LLMs,GenAI

{ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info}/RECORD

@@ -1,34 +1,34 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=rMTM89YDPkmLKlUQV_Rjr7B2KpcvldAHzfpgUqaOqdU,12216
+ai_edge_quantizer/algorithm_manager.py,sha256=UZVS6ZClIAyaX9RzhXvvymbQv_scR0ybMPYl2CgSPVo,12346
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=w6bSONvXkX6bzXAGc0-7b6gNDt9oz9ieq97KP8Sg_JU,7666
 ai_edge_quantizer/calibrator.py,sha256=Sms7_AIHPH9G5xFaz5Ef3a5gPhxuIWQI8d2LUM8C96I,12071
 ai_edge_quantizer/calibrator_test.py,sha256=C_oWOaRugPKYX74jF-eRFH-k6nGOdA8I9_uPiocaOuE,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=zghBh9dTB-ouPFumV-0siBSnEbp0WxF6tGOsn3TLirg,11242
+ai_edge_quantizer/default_policy.py,sha256=0Am2TrgyV7gNl7dbul07rVp58OKDuPyJW9SIqRTrD2g,11280
 ai_edge_quantizer/model_modifier.py,sha256=teGa8I6kGvn6TQY6Xv53YFIc_pQEhNvM9Zb4bvhezyw,7110
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=Hj0_5o-Oa3dSlJ3ryVjRhvsyelHNyek1GrtG9buMczg,13153
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
 ai_edge_quantizer/params_generator.py,sha256=gC7G6Ne4Fumc8RSmIAbx96ZBhszZlHqBKSmE9p6RPTo,20099
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=kX1AoD-YlHYbDI1RfGVXIbPn-CYT7HUF2x77-hPtKBM,16565
+ai_edge_quantizer/qtyping.py,sha256=vq-9jwDViSndHhcC1_RVu2Bk0qu5MgYPGLTRO9z2Naw,16604
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
 ai_edge_quantizer/recipe.py,sha256=FR0uJceumZrnle2VRSOQZ1uXup4S1cTYKRH-N53mWRo,2919
 ai_edge_quantizer/recipe_manager.py,sha256=qcGUD7e7BISKdsY9WH2rdaRR3acmzSA5qMezGNbzlpo,8931
 ai_edge_quantizer/recipe_manager_test.py,sha256=GVOfGFZPRciUb4EF4GkSi6d96LdjS6PbUkAJ0ayy0k8,32243
 ai_edge_quantizer/recipe_test.py,sha256=Fg_sfxovI2fRjk5qdu18ghOvXdUvhDR1TxbE0GHDczc,3381
-ai_edge_quantizer/transformation_instruction_generator.py,sha256=B_TQQe9_Qs7UKXLjMMuz5lORUvXyZOxBS2SpntTnkI8,28077
-ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=E0QSDCav6N6izlJ-a1ZJOsb2VEUxuxBmTbt0-EgDdxY,49890
-ai_edge_quantizer/transformation_performer.py,sha256=nkkqbs81ITB5u2FoWeG9z5d8EtLtCiltOxcQ34okN8E,13091
+ai_edge_quantizer/transformation_instruction_generator.py,sha256=iMGXy7_ufqgQRzu4drAfO31VGdze35peEFh1BMZlVHk,27714
+ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=Zw3EOSnvzjuB4NWeo129eJZxK_EHno9oF9OtEQ-0dnM,48905
+ai_edge_quantizer/transformation_performer.py,sha256=o4J6OUbI0dLoobVYjkOFw5Po3yH0gZJXrfuTIYais4o,13029
 ai_edge_quantizer/transformation_performer_test.py,sha256=xk6A3LStCyPclN51--9uO7XjSxNfZmpdfvrzOL0maNM,20349
 ai_edge_quantizer/algorithms/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting.py,sha256=Bs9CK7wZAw6jNaZ8xEtbwO2vM34VYXNZSMVWvxJo9nw,9297
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py,sha256=EqIHGEZ1LgUrTN7zf880RuAzEv3Qy7kgh5ivObJGHSo,22646
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=rImKK2ax7LrRx6XurSdvRTk0h6WtFGtQn9sYNJcn-uw,30222
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=juXQWnzBTvasxydeDTNxWuE9ag9j6GOmfHMjC4JQu1Y,30799
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=GGf_n3wIeg3GB_eGsmyNJ0fTcxgpeMMbugTMRONK6TQ,3553
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=BDdn_uBZakfHyzdMJPKadsOqxqyC-s6W2ZzFH99L4fE,8652
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
@@ -41,7 +41,7 @@ ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=sha1d99Xk87bI
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=W2QbXP96xeleAmA7qFwco1iq_bOtArGDK6Qj_g6kNl8,15986
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=MgG7Qh2_z4I6InBqEEDSVlaR0q48aMz4xqAlxeG2EMk,12436
 ai_edge_quantizer/algorithms/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/utils/common_utils.py,sha256=UoZxeAQmZk3b3hK51KFwq6XfdbeduXVjdYIxAxlAzB8,34982
+ai_edge_quantizer/algorithms/utils/common_utils.py,sha256=QrEeCuvA7gY_vK1nbKtqassNDClyAjN1ClZIiw63k5U,35895
 ai_edge_quantizer/algorithms/utils/common_utils_test.py,sha256=zqapGEfYhjQWe9cNGPLmdbwtEUUYQRhlO_kNe0cXX6E,18104
 ai_edge_quantizer/transformations/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/transformations/dequant_insert.py,sha256=sL1LHFVzBDSd9jgrzlHz38LWU0bwmVX7iBkaNcui0ts,3566
@@ -50,28 +50,26 @@ ai_edge_quantizer/transformations/duplicate_buffer.py,sha256=TvTHbm24IiICNkWOlvR
 ai_edge_quantizer/transformations/duplicate_buffer_test.py,sha256=YYWl3Q5WF60s8T8pLzzA8TCSxz-i7dqc03dJt1LtMw4,3880
 ai_edge_quantizer/transformations/duplicate_tensor.py,sha256=WKhf2LIAL0MnZe88b6942A37lvHXe1cFjUDqE5VNmvU,2490
 ai_edge_quantizer/transformations/duplicate_tensor_test.py,sha256=s-RqSxNBMfVJyCunXz2eb7-KA6UiBmbOmL7phLslENQ,5056
-ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xeMv3OIymukUy_yW1zK0xN8Ann6I4,13602
-ai_edge_quantizer/transformations/emulated_subchannel_test.py,sha256=gZP6u9NdPXl7s19qB_Un8evou9ZZV6I9Gy0E1rdobHM,7722
 ai_edge_quantizer/transformations/insert_hadamard_rotation.py,sha256=rBbKgcVKHie38NT2UQ7KQ1xCb2tRu_rVl0yFloOAW_A,7562
 ai_edge_quantizer/transformations/insert_hadamard_rotation_test.py,sha256=iV1p3nZfHUATV2YRoBOYurnu3pLy8n3aFppLWGQOPdA,7268
 ai_edge_quantizer/transformations/quant_insert.py,sha256=jn6HsJaV-sqBiFPY-Aqbd64t8zgcYVkEkZI375x_FWY,3958
 ai_edge_quantizer/transformations/quant_insert_test.py,sha256=X9ptPDvJCFkR5tejKnD1SlHFGPazQTW-wNNMV9MEAuw,10107
 ai_edge_quantizer/transformations/quantize_tensor.py,sha256=kjaNrw9mnrn0t8u0vey9S_uPz3iVUicwy4rluxVqV3E,7617
 ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-ZN8ADn5tBBJlqjTWa7ZUN8Mmu5Rcw,9116
-ai_edge_quantizer/transformations/transformation_utils.py,sha256=GwIaKVsePZYgVG2lSanOswcaZYMjvgyqstDVwXl9DGY,6923
+ai_edge_quantizer/transformations/transformation_utils.py,sha256=efJdAkA24wlg6Vj5NFO7_7MDuvQLSNn-l11Vs_JPktI,7123
 ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=MWgq29t7rvxRQIfi4ny9IoODFCTcbpjnIwoCL40zDKk,8698
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/utils/calibration_utils.py,sha256=e3dG7Nm94Ix0hkTWTWPUhEG6a8QR_cAM3PSwblfJV5g,15106
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=4BlksXl7b4yptL8xPR67hmJCnjhN9V10a2PunzfHrUE,9372
 ai_edge_quantizer/utils/test_utils.py,sha256=spqUmSNciOKPQHCBkHE7Zo34eMFq_BfBCAnMT3jAulU,8615
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=pZv8FMWyjBSLN5MGJ2K_dZ6oqkJGbp9RI4CfnlPuPII,10830
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=NnD57Gkx9upNP8Mso-_yp8Z3x1AqlIWb06jPg-hyRkc,10890
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=vTyy6-4PgfFPL3C8uTq_iPFBwdxCjhrWzUiec4DdFPw,14323
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=6fjkM-rycZ95L4yfvlr0TN6RlrhfPzxNUYrZaYO_F0A,12013
 ai_edge_quantizer/utils/validation_utils.py,sha256=oYw33Sg547AqtGw-choPUJmp9SAKkV46J_ddqSsum2Q,3950
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=V_qNDikPD4OPB-siOLQCWNVWTAu87h2IgNYt7teFd-o,2934
-ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info/METADATA,sha256=0-WpgPHWtwW_Wvysp7yPgMXb6nNP6sXI-vJphPlKrBs,1528
-ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.3.0.dev20250622.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info/METADATA,sha256=aigmp6Hzdxwsj0hwX5ARfya0brbvfrjYq-nMdlcQja4,1528
+ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.3.0.dev20250624.dist-info/RECORD,,

ai_edge_quantizer/transformations/emulated_subchannel.py

@@ -1,363 +0,0 @@
-# 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.
-# ==============================================================================
-
-"""Transformation pattern for emulated subchannel quantization."""
-
-from typing import cast
-import numpy as np
-from ai_edge_quantizer import qtyping
-from ai_edge_quantizer.transformations import quantize_tensor
-from ai_edge_quantizer.transformations import transformation_utils
-from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
-
-
-def emulated_subchannel(
-    transformation_input: transformation_utils.TransformationInput,
-) -> qtyping.TransformationInfo:
-  """Emulated subchannel quantization for fully_connected op.
-
-  The input tensor must also be the weight tensor of the fully_connected op.
-
-  after the transformation, the fully connected op will be replaced by:
-  reshape -> batch_matmul -> mul -> sum -> add (if bias is present) ->
-  activation (if fused activation function exist, only support ReLU for now)
-
-  Args:
-    transformation_input: input structure that contains all information needed
-      for the transformation.
-
-  Returns:
-    The transformation info.
-  """
-  # only apply to a single fully_connected op
-  if len(transformation_input.consumers) > 1:
-    raise ValueError('Emulated Subchannel transformation only support one op')
-  if isinstance(
-      transformation_input.quant_params, qtyping.NonLinearQuantParams
-  ):
-    raise ValueError(
-        'Emulated Subchannel transformation only support uniform quantization'
-    )
-  if (
-      transformation_input.op_codes[
-          transformation_input.subgraph.operators[
-              transformation_input.consumers[0]
-          ].opcodeIndex
-      ].builtinCode
-      != schema_py_generated.BuiltinOperator.FULLY_CONNECTED
-  ):
-    raise ValueError(
-        'Emulated Subchannel transformation only support fully_connected op'
-    )
-  if transformation_input.producer != -1:
-    raise ValueError(
-        'Emulated Subchannel transformation only support constant tensor'
-    )
-
-  # insert all tne necessary op codes into the model
-  reshape_op_code_idx = transformation_utils.add_op_code(
-      schema_py_generated.BuiltinOperator.RESHAPE, transformation_input.op_codes
-  )
-  bmm_op_code_idx = transformation_utils.add_op_code(
-      schema_py_generated.BuiltinOperator.BATCH_MATMUL,
-      transformation_input.op_codes,
-  )
-  mul_op_code_idx = transformation_utils.add_op_code(
-      schema_py_generated.BuiltinOperator.MUL, transformation_input.op_codes
-  )
-  sum_op_code_idx = transformation_utils.add_op_code(
-      schema_py_generated.BuiltinOperator.SUM, transformation_input.op_codes
-  )
-
-  original_fc_op_idx = transformation_input.consumers[0]
-  if cast(
-      schema_py_generated.FullyConnectedOptionsT,
-      transformation_input.subgraph.operators[
-          original_fc_op_idx
-      ].builtinOptions,
-  ).fusedActivationFunction not in (
-      schema_py_generated.ActivationFunctionType.NONE,
-      schema_py_generated.ActivationFunctionType.RELU,
-  ):
-    raise ValueError(
-        'Emulated Subchannel transformation only support'
-        ' fusedActivationFunction NONE and RELU for now'
-    )
-
-  weight_tensor = transformation_input.subgraph.tensors[
-      transformation_input.tensor_id
-  ]
-  weight_tensor.type = quantize_tensor.quant_params_to_tflite_type(
-      transformation_input.quant_params.num_bits
-  )
-
-  # modify the weight tensor with the correct quantization parameters
-  transformation_input.buffers[weight_tensor.buffer].data = np.frombuffer(
-      cast(
-          np.ndarray, transformation_input.quant_params.quantized_data
-      ).tobytes(),
-      dtype=np.uint8,
-  )
-  weight_tensor.shape = cast(
-      np.ndarray, transformation_input.quant_params.quantized_data
-  ).shape
-  weight_tensor.quantization.scale = np.ones(shape=[1], dtype=np.float32)
-  weight_tensor.quantization.zeroPoint = np.zeros(
-      shape=[1], dtype=np.int64
-  ).flatten()
-
-  # assuming zero point is 0, so no need to add a zero point tensor
-  for val in transformation_input.quant_params.zero_point.flatten():
-    if val != 0:
-      raise ValueError(
-          'Emulated Subchannel transformation only support zero point 0 for now'
-      )
-
-  scale_tensor_id = transformation_utils.add_new_constant_tensor(
-      weight_tensor.name + b'_scale',
-      transformation_input.quant_params.scale,
-      schema_py_generated.TensorType.FLOAT32,
-      transformation_input.subgraph,
-      transformation_input.buffers,
-  )
-
-  # for fully connected op, the reduce axis is always 1
-  reduce_axes_data = np.array([1], dtype=np.int32)
-  reduce_axes_tensor_id = transformation_utils.add_new_constant_tensor(
-      weight_tensor.name + b'_reduce_axes',
-      reduce_axes_data,
-      schema_py_generated.TensorType.INT32,
-      transformation_input.subgraph,
-      transformation_input.buffers,
-  )
-
-  # find the input and output tensor of the fully connected op
-  activation_input_id = transformation_input.subgraph.operators[
-      transformation_input.consumers[0]
-  ].inputs[0]
-  activation_output_id = transformation_input.subgraph.operators[
-      transformation_input.consumers[0]
-  ].outputs[0]
-  activation_input = transformation_input.subgraph.tensors[activation_input_id]
-  activation_output = transformation_input.subgraph.tensors[
-      activation_output_id
-  ]
-
-  if len(activation_input.shape) != 3:
-    raise ValueError(
-        'Emulated Subchannel transformation only support 3D input tensor'
-    )
-  bmm_input_shape = [
-      activation_input.shape[0] * activation_input.shape[1],
-      weight_tensor.shape[1],
-      1,
-      weight_tensor.shape[2],
-  ]
-  intermediate_tensor_shape = [
-      activation_input.shape[0] * activation_input.shape[1],
-      weight_tensor.shape[1],
-      1,
-      weight_tensor.shape[3],
-  ]
-  sum_output_shape = [
-      activation_input.shape[0] * activation_input.shape[1],
-      1,
-      1,
-      weight_tensor.shape[3],
-  ]
-
-  # create constant tensors for reshape
-  reshape1_shape_id = transformation_utils.add_new_constant_tensor(
-      activation_output.name + b'_reshape_op1_shape',
-      np.array(bmm_input_shape, dtype=np.int32),
-      schema_py_generated.TensorType.INT32,
-      transformation_input.subgraph,
-      transformation_input.buffers,
-  )
-  reshape2_shape_id = transformation_utils.add_new_constant_tensor(
-      activation_output.name + b'_reshape_op2_shape',
-      np.array(activation_output.shape, dtype=np.int32),
-      schema_py_generated.TensorType.INT32,
-      transformation_input.subgraph,
-      transformation_input.buffers,
-  )
-
-  # create all intermediate tensors
-  bmm_input_id = transformation_utils.add_new_activation_tensor(
-      activation_output.name + b'_bmm_input',
-      bmm_input_shape,
-      schema_py_generated.TensorType.FLOAT32,
-      transformation_input.subgraph,
-  )
-  mul_input_id = transformation_utils.add_new_activation_tensor(
-      activation_output.name + b'_mul_input',
-      intermediate_tensor_shape,
-      schema_py_generated.TensorType.FLOAT32,
-      transformation_input.subgraph,
-  )
-  sum_input_id = transformation_utils.add_new_activation_tensor(
-      activation_output.name + b'_reduce_sum_input',
-      intermediate_tensor_shape,
-      schema_py_generated.TensorType.FLOAT32,
-      transformation_input.subgraph,
-  )
-  reshape_op2_input_id = transformation_utils.add_new_activation_tensor(
-      activation_output.name + b'_reshape_op2_input',
-      sum_output_shape,
-      schema_py_generated.TensorType.FLOAT32,
-      transformation_input.subgraph,
-  )
-
-  # reshape
-  reshape_op1 = schema_py_generated.OperatorT()
-  reshape_op1.opcodeIndex = reshape_op_code_idx
-  reshape_op1_option = schema_py_generated.ReshapeOptionsT()
-  reshape_op1_option.newShape = bmm_input_shape
-  reshape_op1.inputs = [activation_input_id, reshape1_shape_id]
-  reshape_op1.outputs = [bmm_input_id]
-  reshape_op1.builtinOptionsType = (
-      schema_py_generated.BuiltinOptions.ReshapeOptions
-  )  # reshape option index
-  reshape_op1.builtinOptions = reshape_op1_option
-
-  # batch_matmul
-  bmm_op = schema_py_generated.OperatorT()
-  bmm_op.opcodeIndex = bmm_op_code_idx
-  bmm_op.inputs = [bmm_input_id, transformation_input.tensor_id]
-  bmm_op.outputs = [mul_input_id]
-  bmm_op.builtinOptionsType = (
-      schema_py_generated.BuiltinOptions.BatchMatMulOptions
-  )
-  bmm_op.builtinOptions = schema_py_generated.BatchMatMulOptionsT()
-
-  # mul
-  mul_op = schema_py_generated.OperatorT()
-  mul_op.opcodeIndex = mul_op_code_idx
-  mul_option = schema_py_generated.MulOptionsT()
-  mul_option.fusedActivationFunction = (
-      schema_py_generated.ActivationFunctionType.NONE
-  )
-  mul_op.inputs = [mul_input_id, scale_tensor_id]
-  mul_op.outputs = [sum_input_id]
-  mul_op.builtinOptionsType = schema_py_generated.BuiltinOptions.MulOptions
-  mul_op.builtinOptions = mul_option
-
-  # sum
-  sum_op = schema_py_generated.OperatorT()
-  sum_op.opcodeIndex = sum_op_code_idx
-  sum_op.inputs = [sum_input_id, reduce_axes_tensor_id]
-  sum_op.outputs = [reshape_op2_input_id]
-  sum_op.builtinOptionsType = schema_py_generated.BuiltinOptions.ReducerOptions
-  sum_op.builtinOptions = schema_py_generated.ReducerOptionsT()
-  sum_op.builtinOptions.keepDims = True
-
-  # reshape
-  reshape_op2 = schema_py_generated.OperatorT()
-  reshape_op2.opcodeIndex = reshape_op_code_idx
-  reshape_op2_option = schema_py_generated.ReshapeOptionsT()
-  reshape_op2_option.newShape = activation_output.shape
-  reshape_op2.inputs = [reshape_op2_input_id, reshape2_shape_id]
-  reshape_op2.outputs = [activation_output_id]
-  reshape_op2.builtinOptionsType = (
-      schema_py_generated.BuiltinOptions.ReshapeOptions
-  )
-  reshape_op2.builtinOptions = reshape_op2_option
-
-  transformation_input.subgraph.operators.insert(
-      original_fc_op_idx, reshape_op1
-  )
-  transformation_input.subgraph.operators.insert(original_fc_op_idx + 1, bmm_op)
-  transformation_input.subgraph.operators.insert(original_fc_op_idx + 2, mul_op)
-  transformation_input.subgraph.operators.insert(original_fc_op_idx + 3, sum_op)
-  transformation_input.subgraph.operators.insert(
-      original_fc_op_idx + 4, reshape_op2
-  )
-  ops_added = 5
-  last_op = reshape_op2
-
-  # If there is a bias tensor (the third input to the original fc op),
-  # we need an add to process it. The current fc op id need to be recalculated
-  # because we added operators in front of it.
-  current_fc_op_id = original_fc_op_idx + ops_added
-  if (
-      len(transformation_input.subgraph.operators[current_fc_op_id].inputs) > 2
-      and transformation_input.subgraph.operators[current_fc_op_id].inputs[2]
-      != -1
-  ):
-    add_op_code_idx = transformation_utils.add_op_code(
-        schema_py_generated.BuiltinOperator.ADD, transformation_input.op_codes
-    )
-    reshape_op2_output_id = transformation_utils.add_new_activation_tensor(
-        activation_output.name + b'_reshape_op2_output',
-        activation_output.shape,
-        schema_py_generated.TensorType.FLOAT32,
-        transformation_input.subgraph,
-    )
-    last_op.outputs = [reshape_op2_output_id]
-    add_op = schema_py_generated.OperatorT()
-    add_op.opcodeIndex = add_op_code_idx
-    add_option = schema_py_generated.AddOptionsT()
-    add_op.builtinOptionsType = schema_py_generated.BuiltinOptions.AddOptions
-    add_op.builtinOptions = add_option
-    add_op.inputs = [
-        reshape_op2_output_id,
-        transformation_input.subgraph.operators[
-            original_fc_op_idx + ops_added
-        ].inputs[2],
-    ]
-    add_op.outputs = [activation_output_id]
-    transformation_input.subgraph.operators.insert(
-        original_fc_op_idx + ops_added, add_op
-    )
-    ops_added += 1
-    last_op = add_op
-
-  # If the fused activation function is RELU, we need to add a relu op.
-  # The current fc op id need to be recalculated because we added operators
-  # in front of it.
-  fc_fused_activation_function = cast(
-      schema_py_generated.FullyConnectedOptionsT,
-      transformation_input.subgraph.operators[
-          original_fc_op_idx + ops_added
-      ].builtinOptions,
-  ).fusedActivationFunction
-  if (
-      fc_fused_activation_function
-      == schema_py_generated.ActivationFunctionType.RELU
-  ):
-    activation_output.name += b'_relu'
-    relu_input_id = transformation_utils.add_new_activation_tensor(
-        activation_output.name + b'_relu_input',
-        activation_output.shape,
-        schema_py_generated.TensorType.FLOAT32,
-        transformation_input.subgraph,
-    )
-    last_op.outputs = [relu_input_id]
-    relu_op = schema_py_generated.OperatorT()
-    relu_op.opcodeIndex = transformation_utils.add_op_code(
-        schema_py_generated.BuiltinOperator.RELU, transformation_input.op_codes
-    )
-    relu_op.inputs = [relu_input_id]
-    relu_op.outputs = [activation_output_id]
-    transformation_input.subgraph.operators.insert(
-        original_fc_op_idx + ops_added, relu_op
-    )
-    ops_added += 1
-    last_op = relu_op
-
-  del transformation_input.subgraph.operators[original_fc_op_idx + ops_added]
-  return qtyping.TransformationInfo(
-      original_fc_op_idx, ops_added - 1, activation_output_id
-  )

ai_edge_quantizer/transformations/emulated_subchannel_test.py

@@ -1,212 +0,0 @@
-# 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.
-# ==============================================================================
-
-"""Tests for emulated_subchannel."""
-
-import os
-import numpy as np
-from tensorflow.python.platform import googletest
-from ai_edge_quantizer import qtyping
-from ai_edge_quantizer.transformations import emulated_subchannel
-from ai_edge_quantizer.transformations import transformation_utils
-from ai_edge_quantizer.utils import test_utils
-from ai_edge_quantizer.utils import tfl_flatbuffer_utils
-from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
-
-TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("..")
-
-
-class EmulatedSubchannelTest(googletest.TestCase):
-  """Tests for emulated_subchannel."""
-
-  def setUp(self):
-    super().setUp()
-    self.params = qtyping.UniformQuantParams(
-        num_bits=8,
-        quantized_dimension=None,
-        scale=np.ones([1, 1, 1, 4], dtype=np.float32),
-        zero_point=np.zeros([1, 1, 1, 4], dtype=np.int64),
-        symmetric=True,
-        quantized_data=np.ones([1, 4, 2, 4], dtype=np.int8),
-    )
-
-  def test_emulate_subchannel_without_bias_succeeds(self):
-    """Tests the emulated_subchannel function."""
-    self._model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_no_bias.tflite"
-    )
-    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
-    subgraph = self._model.subgraphs[0]
-    model = self._model
-    ret = emulated_subchannel.emulated_subchannel(
-        transformation_utils.TransformationInput(
-            tensor_id=1,
-            op_codes=model.operatorCodes,
-            buffers=model.buffers,
-            subgraph=subgraph,
-            producer=-1,
-            consumers=[0],
-            quant_params=self.params,
-        )
-    )
-    self.assertEqual(ret.op_id, 0)
-    self.assertEqual(ret.num_ops_added, 4)
-    self.assertEqual(ret.output_tensor_id, 2)
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.RESHAPE,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.BATCH_MATMUL,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[2].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.MUL,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[3].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.SUM,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[4].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.RESHAPE,
-    )
-    self.assertEqual(
-        subgraph.tensors[subgraph.operators[2].inputs[1]].name,
-        b"arith.constant_scale",
-    )
-    self.assertListEqual(
-        np.frombuffer(
-            model.buffers[
-                subgraph.tensors[subgraph.operators[2].inputs[1]].buffer
-            ].data,
-            dtype=np.float32,
-        ).tolist(),
-        np.ones([1, 1, 1, 4]).flatten().tolist(),
-    )
-
-  def test_emulate_subchannel_with_bias_succeeds(self):
-    """Tests the emulated_subchannel function."""
-    self._model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
-    )
-    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
-    subgraph = self._model.subgraphs[0]
-    model = self._model
-    ret = emulated_subchannel.emulated_subchannel(
-        transformation_utils.TransformationInput(
-            tensor_id=1,
-            op_codes=model.operatorCodes,
-            buffers=model.buffers,
-            subgraph=subgraph,
-            producer=-1,
-            consumers=[0],
-            quant_params=self.params,
-        )
-    )
-    self.assertEqual(ret.op_id, 0)
-    self.assertEqual(ret.num_ops_added, 5)
-    self.assertEqual(ret.output_tensor_id, 3)
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.RESHAPE,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.BATCH_MATMUL,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[2].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.MUL,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[3].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.SUM,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[4].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.RESHAPE,
-    )
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[5].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.ADD,
-    )
-    self.assertEqual(
-        subgraph.tensors[subgraph.operators[2].inputs[1]].name,
-        b"arith.constant_scale",
-    )
-    self.assertListEqual(
-        np.frombuffer(
-            model.buffers[
-                subgraph.tensors[subgraph.operators[2].inputs[1]].buffer
-            ].data,
-            dtype=np.float32,
-        ).tolist(),
-        np.ones([1, 1, 1, 4]).flatten().tolist(),
-    )
-
-  def test_emulated_subchannel_with_fused_relu_succeeds(self):
-    """Tests the emulated_subchannel function with fused relu."""
-    self._model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias_relu.tflite"
-    )
-    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
-    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
-    subgraph = self._model.subgraphs[0]
-    model = self._model
-    ret = emulated_subchannel.emulated_subchannel(
-        transformation_utils.TransformationInput(
-            tensor_id=1,
-            op_codes=model.operatorCodes,
-            buffers=model.buffers,
-            subgraph=subgraph,
-            producer=-1,
-            consumers=[0],
-            quant_params=self.params,
-        )
-    )
-    self.assertEqual(ret.op_id, 0)
-    self.assertEqual(ret.num_ops_added, 6)
-    self.assertEqual(ret.output_tensor_id, 3)
-    self.assertEqual(
-        model.operatorCodes[subgraph.operators[6].opcodeIndex].builtinCode,
-        schema_py_generated.BuiltinOperator.RELU,
-    )
-
-  def test_emulated_subchannel_raises_when_unsupported_activation(self):
-    """Tests the emulated_subchannel function with unsupported activation."""
-    self._model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias_relu6.tflite"
-    )
-    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
-    subgraph = self._model.subgraphs[0]
-    model = self._model
-    with self.assertRaises(ValueError):
-      emulated_subchannel.emulated_subchannel(
-          transformation_utils.TransformationInput(
-              tensor_id=1,
-              op_codes=model.operatorCodes,
-              buffers=model.buffers,
-              subgraph=subgraph,
-              producer=-1,
-              consumers=[0],
-              quant_params=self.params,
-          )
-      )
-
-
-if __name__ == "__main__":
-  googletest.main()