Create a deep copy of the model while preserving specific attributes. Args: model (torch.nn.Module): The model to deep copy. Returns

Get the dtype of the node.

Takes one float input to an op, and converts it to quantized form.

Fuse the Requantize/Dequantize with fused quantized Ops.

Get all the layers that could be smooth quantized. Returns: list: All the smooth quantization layer names.

Run the calibration on the input graph. Args: model(TensorflowBaseModel): input TensorflowBaseModel

Converts to a jit model. Args: model (QuantizationInterceptionModule): a prepared model. example_inputs (tensor/tuple/dict):

Parse the graph and get the input node and output node name details.

Core function used for element-wise quantization. Args: A (torch.Tensor): tensor to be quantized bits (int): number of mantissa b

Use absolute paths from manifest directly; fallback to root for relative paths.

Scale the layer weights at input channel and depthwise convolution output channel. Args: layer_name (str): The layer name.

Identify the layers which can be smooth quantized. Args: op_types (list): The operation types to be smooth quantized.

Adds a trial result to the tuning history. Args: trial_index (int): The index of the trial. trial_result (Union[int,

Dataset return data per batch.

Collect output of block. Args: model (torch.nn.Module): input model. Returns: output(list): a list of block

Build a dataloader that captures input data and keyword arguments used in forward passes of the model. Args: model (torch.nn.Module):

Calculate the scaling factor for weights based on the input max values and weight magnitudes. Args: input_max_abs (Tensor): The maximu

Process the absorb layer and smooth layers, then return the channel-wise max value info. Args: absorb_to_layer (dict): A dicti

Default call before add_variables is invoked. Should not be reached after setup.

Check configs and quantization configs. Args: user_cfg (dict): quantization configuration for ops. cfgs (dict): configs loade

Check if the current Keras backend is JAX and log a warning or error if not.

Collate the prediction results.

Converts a prepared model to a quantized model. Args: model (QuantizationInterceptionModule): the prepared model to be convert

Converts `text` to Unicode (if it's not already), assuming utf-8 input.

Return the count of FusedBatchNorm op and the count of primitive ops which may make up batchnorm computation in a given graph.

Create testing graph and compute the result from original graph. Args: pattern_match_mode: A label string to indicate which batchn

Create constant node. Args: name (string): op name value (np.array): input data dtype (datatype): data ty

Dequantize FP8 weights using the module hook when required. Args: mod (torch.nn.Module): Module that may provide a dequantization callbac

Fold all the sequences only consist of const and self.supported_op_type. Args: input_graph_def (graphdef): graphdef object

This is a function to dump quantizable ops of model to user. Args: user_cfg (dict): quantization config. Returns: None

Log the end of the evaluation process.

Log the start of the evaluation process.

Log the end of the execution process.

Log the start of the execution process.

Exports a eager model for PT2E quantization. Args: model (torch.nn.Module): The PyTorch model to be exported. example_inputs (Tup

Output log with the fatal level.

Create a StaticQuantConfig from a dictionary. Args: config_dict (Dict): Configuration fields. Returns: Stati

Get element format with str.

Fuse the linear operations in Layernorm into the adjacent linear blocks.

To generate the scheme requested by vllm compatible config.

Get all registered configuration classes for a specific framework. Args: fwk_name (str): The framework name. Returns:

Retrieves all the configuration sets from the config registry for a given framework name. Args: fwk_name (str): The name of the framework

Returns the best trial record based on the trial result. Raises: AssertionError: If there are no trial records in the tuning moni

Get the block names for transformers-like networks. Args: model: The model. Returns: block_names: A list whose elements are list of

Get prefix and number of attention blocks of transformer models. Args: model (torch.nn.Module): input model Returns: block_p

Serialize this class to a config dict.

Serialize this class to a config dict.

Generate a DataLoader for calibration using specified parameters. Args: tokenizer (Tokenizer): The tokenizer to use for tokenization.

Get the default configuration of double quant. Args: type (str, optional): double quant type. Defaults to "BNB_NF4". Returns:

Generate the default static quant config. Returns: the default keras config.

Convert a string to the corresponding enum value of enum_class. Raises ValueError if the string does not match any enum member.

Get supported patterns by low precisions. Returns: [dictionary list]: A list composed of dictionary which key is precision

Get the input shape of the input model.

Get module from model by key name. Args: module (torch.nn.Module): original model key (str): module name to be replaced

Only necessary for the newly registered patched module that doesn't in _mod_default_dict. Return the module info for the module, which is used

Returns the number of evaluation functions in the eval_fn_registry. Returns: int: The number of evaluation functions.

Get the supported op types by all precisions. Returns: [dictionary list]: A list composed of dictionary which key is precision

Executed the non-precision dependent graph optimization. The input graph will be optimized with following passes: 1. Remove the train

Get a dictionary containing the parameters and their values for the current instance. Returns: A dictionary containing the parame

Initializes the tuning process. Args: tuning_config (TuningConfig): The configuration for the tuning process. Returns: Tuple

Inverts the scale of the input tensor, list of tensors, or tuple of tensors.

Dynamically checks if the given op supports dynamic quantization by looking up its ModuleInfo and checking for a 'supports_dynamic_quantizati

Load(f, map_location=None, pickle_module=pickle, *, weights_only=False, **pickle_load_args). Loads an object saved with :func:`torch.save` from a

This is a function to load TF graph from pb file Args: model_file (string): TF pb file local path Returns: graph: TF graph o

Perfoms a matmul with lp inputs and returns output in hp

Constructs a scale_method config dict for fp8_config, with a default and an override section, and also returns the expected_configs dict for test

The model forward function.

Auto mapping input to device for all kinds of format. Args: input (torch.tensor): input data device (str, optional): target devic

Check whether to apply this algorithm according to configs_mapping. Args: configs_mapping (Dict[Tuple[str, callable], BaseConfig]): confi

Check if need to stop tuning. Either accuracy goal is met, max trials is reached or timeout is reached. Returns: stop_flag: True