ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -0,0 +1,255 @@
+# Copyright 2024 The AI Edge Torch 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.
+# ==============================================================================
+
+import dataclasses
+import tempfile
+from typing import Any, Dict, List, Optional, Tuple
+
+from ai_edge_torch import odml_torch
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.lowertools import translate_recipe
+from ai_edge_torch.odml_torch import export
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+import torch
+
+from tensorflow.compiler.tf2xla.python import xla as tfxla
+from tensorflow.lite.python import conversion_metadata_schema_py_generated as conversion_metadata_fb
+
+MlirBundle = odml_torch.export.MlirLowered
+
+
+@dataclasses.dataclass
+class MergedBundle:
+  """A bundle of MlirLowered that has been merged."""
+
+  bundles: list[odml_torch.export.MlirLowered]
+  exported_programs: list[torch.export.ExportedProgram]
+  deduped_tf_vars: list[tf.Variable]
+
+
+def torch_dtype_to_tf(dtype):
+  return {
+      torch.double: tf.float64,
+      torch.float32: tf.float32,
+      torch.half: tf.float16,
+      torch.long: tf.int64,
+      torch.int32: tf.int32,
+      torch.int16: tf.int16,
+      torch.bool: tf.bool,
+  }.get(dtype)
+
+
+def _get_shape_with_dynamic(signature: export.VariableSignature):
+  return [
+      None if export_utils.is_torch_dynamic(s) else s for s in signature.shape
+  ]
+
+
+def _extract_call_args(
+    bundle: export.MlirLowered,
+    args: Tuple[Any],
+    tf_state_dict: Dict[str, tf.Variable],
+):
+  call_args = []
+  for sig in bundle.input_signature:
+    if sig.input_spec.is_user_input:
+      call_args.append(args[sig.input_spec.i])
+    elif sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      call_args.append(tf_state_dict[name])
+  return call_args
+
+
+def _wrap_as_tf_func(
+    bundle: export.MlirLowered,
+    tf_state_dict: Dict[str, tf.Variable],
+    exported_program: torch.export.ExportedProgram,
+):
+  def inner(*args):
+    t_outs = [torch_dtype_to_tf(sig.dtype) for sig in bundle.output_signature]
+    s_outs = [_get_shape_with_dynamic(sig) for sig in bundle.output_signature]
+    call_args = _extract_call_args(bundle, args, tf_state_dict)
+    # HACK: In OSS, we use MLIR pybinding and StableHLO dialect from JAX's
+    # build, which may not have the same StableHLO version as what used in
+    # TFLite converter. Therefore we always serialize MLIR module in VHLO.
+    # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
+    call_module_return = tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=t_outs,  # dtype information
+        Sout=s_outs,  # Shape information
+        function_list=[],
+        module=bundle.module_bytecode_vhlo,
+    )
+    spec = exported_program.call_spec.out_spec
+
+    # The module returning a flat array.
+    if not spec.context:
+      return call_module_return
+
+    flat_names = common_utils.flat_dict_names(spec.children_specs, spec.context)
+    return {name: value for name, value in zip(flat_names, call_module_return)}
+
+  return inner
+
+
+def _make_tf_signature(
+    input_signature: list[export.VariableSignature],
+    signature: signature_module.Signature,
+) -> List[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  user_input_signature = sorted(
+      [sig for sig in input_signature if sig.input_spec.is_user_input],
+      key=lambda sig: sig.input_spec.i,
+  )
+  tf_signature = []
+
+  for sig in user_input_signature:
+    shape = _get_shape_with_dynamic(sig)
+    tf_signature.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=torch_dtype_to_tf(sig.dtype),
+            name=input_names[sig.input_spec.i],
+        )
+    )
+  return tf_signature
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+):
+  tf_state_dict = merged_bundle.bundles[0].state_dict
+
+  tf_signatures = [
+      _make_tf_signature(bundle.input_signature, sig)
+      for bundle, sig in zip(merged_bundle.bundles, signatures)
+  ]
+  tf_functions = [
+      _wrap_as_tf_func(bundle, tf_state_dict, ep)
+      for bundle, ep in zip(
+          merged_bundle.bundles, merged_bundle.exported_programs
+      )
+  ]
+
+  tf_module = tf.Module()
+  tf_module.f = []
+
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = merged_bundle.deduped_tf_vars
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
+    converter._experimental_enable_composite_direct_lowering = True
+    converter.model_origin_framework = "PYTORCH"
+
+    conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      translated_recipe = translate_recipe.translate_to_ai_edge_recipe(
+          quant_config.generative_recipe
+      )
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      tflite_model = translate_recipe.quantize_model(
+          tflite_model, translated_recipe
+      )
+
+  return tflite_model
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return odml_torch.export.exported_program_to_mlir(exported_program).get_text(
+      enable_debug_info=True
+  )
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
+) -> export.MlirLowered:
+  """Converts a ExportedProgram to a MlirLowered."""
+  return odml_torch.export.exported_program_to_mlir(exported_program)
+
+
+def merge_mlir_bundles(
+    bundles: list[export.MlirLowered],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> MergedBundle:
+  """Merges a list of MlirLowered into one."""
+  state_dict, deduped_vars = common_utils.gather_state_dict(
+      exported_programs, signatures
+  )
+
+  merged_bundle = MergedBundle(
+      bundles=bundles.copy(),
+      exported_programs=exported_programs,
+      deduped_tf_vars=deduped_vars,
+  )
+  for bundle, signature in zip(merged_bundle.bundles, signatures):
+    bundle.state_dict = state_dict
+
+    for var_sig in bundle.input_signature:
+      if var_sig.input_spec.is_parameter:
+        var_sig.input_spec.name = signature.name + "_" + var_sig.input_spec.name
+
+  return merged_bundle

ai_edge_torch/lowertools/test_utils.py

@@ -0,0 +1,60 @@
+# Copyright 2024 The AI Edge Torch 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.
+# ==============================================================================
+
+import re
+from typing import Optional
+from ai_edge_torch import config
+from absl.testing import absltest as googletest
+
+
+def _extract_backend_configs(mlir):
+  mlir = mlir.replace("\\22", '"')
+  configs = []
+  for match in re.finditer(r"backend_config\s*=\s*\"(\{.*\})\"", mlir):
+    configs.append(match.group(1))
+  return "\n".join(configs)
+
+
+def assert_string_count(
+    test_case: googletest.TestCase,
+    mlir: str,
+    torch_xla_pattern_counter: dict[str, int],
+    odml_torch_pattern_counter: dict[str, int],
+    odml_torch_attr_counter: Optional[dict[str, int]] = None,
+):
+
+  if odml_torch_attr_counter is None:
+    odml_torch_attr_counter = {}
+
+  if config.Config.use_torch_xla:
+    for key in torch_xla_pattern_counter:
+      test_case.assertEqual(
+          mlir.count(key),
+          torch_xla_pattern_counter[key],
+      )
+  else:
+    for key in odml_torch_pattern_counter:
+      test_case.assertEqual(
+          mlir.count(key),
+          odml_torch_pattern_counter[key],
+      )
+    backend_configs = _extract_backend_configs(mlir)
+    print("backend_configs:")
+    print(backend_configs)
+    for key in odml_torch_attr_counter:
+      test_case.assertEqual(
+          backend_configs.count(key),
+          odml_torch_attr_counter[key],
+      )

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -0,0 +1,284 @@
+# Copyright 2024 The AI Edge Torch 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.
+# ==============================================================================
+
+import copy
+import dataclasses
+from dataclasses import dataclass
+import gc
+import itertools
+import logging
+import tempfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+from ai_edge_torch import model
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.lowertools import translate_recipe
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+from torch_xla import stablehlo
+
+try:
+  import tensorflow as tf
+
+  from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+  from tensorflow.lite.python import conversion_metadata_schema_py_generated as conversion_metadata_fb  # isort:skip
+except ImportError:
+  logging.error(
+      "This module needs tensorflow with xla support.\n"
+      "Please install tensorflow with `pip install tf-nightly`.\n"
+  )
+  raise
+
+MlirBundle = stablehlo.StableHLOModelBundle
+
+
+@dataclasses.dataclass
+class MergedBundle:
+
+  bundle: stablehlo.StableHLOModelBundle
+  exported_programs: list[torch.export.ExportedProgram]
+  deduped_tf_vars: list[tf.Variable]
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
+) -> stablehlo.StableHLOModelBundle:
+  # Setting export_weights to False here so that pytorch/xla avoids copying the
+  # weights to a numpy array which would lead to memory bloat. This means that
+  # the state_dict in the returned bundle is going to be empty.
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program,
+      stablehlo.StableHLOExportOptions(
+          override_tracing_arguments=sample_args, export_weights=False
+      ),
+  )._bundle
+
+
+def merge_mlir_bundles(
+    bundles: list[stablehlo.StableHLOModelBundle],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> stablehlo.StableHLOGraphModule:
+  state_dict, deduped_tf_vars = common_utils.gather_state_dict(
+      exported_programs, signatures
+  )
+
+  new_shlo_model_bundle = stablehlo.StableHLOModelBundle(
+      state_dict=state_dict, additional_constants=[], stablehlo_funcs=[]
+  )
+
+  for bundle, signature in zip(bundles, signatures):
+    const_offset = len(new_shlo_model_bundle.additional_constants)
+    for func in bundle.stablehlo_funcs:
+      func.meta.name = signature.name + "_" + func.meta.name
+      for loc in func.meta.input_locations:
+        if loc.type_ == stablehlo.VariableType.CONSTANT:
+          loc.position += const_offset
+        elif loc.type_ == stablehlo.VariableType.PARAMETER:
+          loc.name = signature.name + "_" + loc.name
+      new_shlo_model_bundle.stablehlo_funcs.append(func)
+    new_shlo_model_bundle.additional_constants.extend(
+        bundle.additional_constants
+    )
+  return MergedBundle(
+      bundle=new_shlo_model_bundle,
+      exported_programs=exported_programs,
+      deduped_tf_vars=deduped_tf_vars,
+  )
+
+
+def _get_shape_with_dynamic(signature: stablehlo.VariableSignature):
+  shape = copy.copy(signature.shape)
+  for i in signature.dynamic_dims:
+    shape[i] = None
+  return shape
+
+
+def _wrap_as_tf_func(
+    func: stablehlo.StableHLOFunc,
+    bundle: stablehlo.StableHLOModelBundle,
+    exported_program: torch.export.ExportedProgram,
+):
+  def inner(*args):
+    type_info = [sig.dtype for sig in func.meta.output_signature]
+    shape_info = [
+        _get_shape_with_dynamic(sig) for sig in func.meta.output_signature
+    ]
+    call_args = stablehlo._extract_call_parameters(args, func.meta, bundle)
+    call_module_return = tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=type_info,
+        Sout=shape_info,
+        function_list=[],
+        module=func.bytecode,
+    )
+    spec = exported_program.call_spec.out_spec
+
+    # The module returning a flat array.
+    if not spec.context:
+      return call_module_return
+
+    flat_names = common_utils.flat_dict_names(spec.children_specs, spec.context)
+    return {name: value for name, value in zip(flat_names, call_module_return)}
+
+  return inner
+
+
+def _make_tf_signature(
+    meta: stablehlo.StableHLOFunctionMeta,
+    signature: signature_module.Signature,
+) -> list[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  input_pos_to_spec = {
+      loc.position: spec
+      for loc, spec in itertools.chain(
+          zip(meta.input_locations, meta.input_signature), meta.unused_inputs
+      )
+      if loc.type_ == stablehlo.VariableType.INPUT_ARG
+  }
+  assert len(input_pos_to_spec) == len(input_names)
+
+  primitive_type_to_tf_type = {"int": "int32", "float": "float32"}
+  ret: list[tf.TensorSpec] = []
+  for i, name in enumerate(input_names):
+    spec = input_pos_to_spec[i]
+    shape = _get_shape_with_dynamic(spec)
+    ret.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=primitive_type_to_tf_type[spec.dtype]
+            if spec.dtype in primitive_type_to_tf_type
+            else spec.dtype,
+            name=name,
+        )
+    )
+  return ret
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program
+  ).get_stablehlo_text()
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+) -> None:
+  """Converts a StableHLOGraphModule to a tflite model.
+
+  Args: shlo_bundle - model to export and save
+
+    signatures: List of signatures from which names of the signatures is
+    extracted.
+    quant_config: User-defined quantization method and scheme of the model.
+    _tfl_converter_flags: A nested dictionary allowing setting flags for the
+    underlying tflite converter.
+  """
+
+  tf_module = tf.Module()
+
+  shlo_bundle = merged_bundle.bundle
+
+  shlo_bundle.additional_constants = [
+      tf.Variable(v, trainable=False) for v in shlo_bundle.additional_constants
+  ]
+  tf_signatures: list[list[tf.TensorSpec]] = list(
+      _make_tf_signature(func.meta, sig)
+      for func, sig in zip(shlo_bundle.stablehlo_funcs, signatures)
+  )
+
+  tf_functions = [
+      _wrap_as_tf_func(func, shlo_bundle, ep)
+      for func, ep in zip(
+          shlo_bundle.stablehlo_funcs, merged_bundle.exported_programs
+      )
+  ]
+
+  tf_module.f = []
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = (
+      merged_bundle.deduped_tf_vars + shlo_bundle.additional_constants
+  )
+  del shlo_bundle
+  gc.collect()
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+    # Clean up intermediate memory early.
+    del tf_module
+    del tf_concrete_funcs
+    gc.collect()
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
+    converter._experimental_enable_composite_direct_lowering = True
+
+    conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      translated_recipe = translate_recipe.translate_to_ai_edge_recipe(
+          quant_config.generative_recipe
+      )
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      tflite_model = translate_recipe.quantize_model(
+          tflite_model, translated_recipe
+      )
+
+  return tflite_model

ai_edge_torch/{generative/quantize/ai_edge_quantizer_glue → lowertools}/translate_recipe.py

@@ -13,14 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-import json
-
 from ai_edge_quantizer import quantizer
-
 from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import quant_recipe
 
-_OpExecutionMode = quantizer.qtyping.OpExecutionMode
+_ComputePrecision = quantizer.qtyping.ComputePrecision
+_QuantGranularity = quantizer.qtyping.QuantGranularity
 _OpName = quantizer.qtyping.TFLOperationName
 _TensorQuantConfig = quantizer.qtyping.TensorQuantizationConfig
 _OpQuantConfig = quantizer.qtyping.OpQuantizationConfig
@@ -44,26 +42,40 @@ def _get_nbits_from_dtype(dtype: quant_attrs.Dtype) -> int:
   raise ValueError('Unimplemented number of bits')
 
 
-def _get_dtype_from_dtype(dtype: quant_attrs.Dtype) -> quantizer.qtyping.TensorDataType:
+def _get_dtype_from_dtype(
+    dtype: quant_attrs.Dtype,
+) -> quantizer.qtyping.TensorDataType:
   if dtype == quant_attrs.Dtype.FP32 or dtype == quant_attrs.Dtype.FP16:
     return quantizer.qtyping.TensorDataType.FLOAT
   else:
     return quantizer.qtyping.TensorDataType.INT
 
 
-def _get_execution_mode_from_mode(mode: quant_attrs.Mode) -> _OpExecutionMode:
+def _get_compute_precision_from_mode(
+    mode: quant_attrs.Mode,
+) -> _ComputePrecision:
   if mode == quant_attrs.Mode.DYNAMIC_RANGE:
-    return _OpExecutionMode.DRQ
+    return _ComputePrecision.INTEGER
   elif mode == quant_attrs.Mode.WEIGHT_ONLY:
-    return _OpExecutionMode.WEIGHT_ONLY
+    return _ComputePrecision.FLOAT
   raise ValueError('Unimplemented execution mode')
 
 
-def _get_channelwise_from_granularity(granularity: quant_attrs.Granularity) -> bool:
-  if granularity == quant_attrs.Granularity.CHANNELWISE:
-    return True
-  elif granularity == quant_attrs.Granularity.NONE:
+def _get_explicit_dequant_from_mode(mode: quant_attrs.Mode) -> bool:
+  if mode == quant_attrs.Mode.DYNAMIC_RANGE:
     return False
+  elif mode == quant_attrs.Mode.WEIGHT_ONLY:
+    return True
+  raise ValueError('Unimplemented execution mode')
+
+
+def _get_granularity(
+    granularity: quant_attrs.Granularity,
+) -> bool:
+  if granularity == quant_attrs.Granularity.CHANNELWISE:
+    return _QuantGranularity.CHANNELWISE
+  if granularity == quant_attrs.Granularity.NONE:
+    return _QuantGranularity.TENSORWISE
   raise ValueError('Unimplemented granularity')
 
 
@@ -87,10 +99,13 @@ def _set_quant_config(
           weight_tensor_config=_TensorQuantConfig(
               num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
               symmetric=True,
-              channel_wise=_get_channelwise_from_granularity(layer_recipe.granularity),
+              granularity=_get_granularity(layer_recipe.granularity),
               dtype=_get_dtype_from_dtype(layer_recipe.weight_dtype),
           ),
-          execution_mode=_get_execution_mode_from_mode(layer_recipe.mode),
+          compute_precision=_get_compute_precision_from_mode(layer_recipe.mode),
+          explicit_dequantize=_get_explicit_dequant_from_mode(
+              layer_recipe.mode
+          ),
       ),
       algorithm_key=_get_algorithm_key_from_algorithm(layer_recipe.algorithm),
   )