def optimize(
    model,
    optimizer=None,
    dtype=torch.float,
    inplace=False,
    device="cpu",
    quantization_config=None,
    qconfig_summary_file=None,
    low_precision_checkpoint=None,
    sample_inputs=None,
    deployment_mode=True,
    cache_weight_for_large_batch=False,
):
    r"""
    Apply optimizations at Python frontend to the given transformers model (nn.Module).
    This API focus on transformers models, especially for generation tasks inference.

    Well supported model family with full functionalities:
    Llama, MLlama, GPT-J, GPT-Neox, OPT, Falcon, Bloom, CodeGen, Baichuan, ChatGLM, GPTBigCode,
    T5, Mistral, MPT, Mixtral, StableLM, QWen, Git, Llava, Yuan, Phi, Qwen3, Whisper. Maira2, Jamba, DeepSeekV2.

    For the model that is not in the scope of supported model family above, will try to
    apply default ipex.optimize transparently to get benifits (not include quantizations,
    only works for dtypes of torch.bfloat16 and torch.half and torch.float).

    Args:
        model (torch.nn.Module): User model to apply optimizations.
        optimizer (torch.optim.Optimizer): User optimizer to apply optimizations
            on, such as SGD. The default value is ``None``, meaning inference case.
        dtype (torch.dtype): Now it works for ``torch.bfloat16``, ``torch.half`` and ``torch.float``.
            The default value is ``torch.float``. When working with quantization, it means the mixed dtype with quantization.
        inplace (bool): Whether to perform inplace optimization. Default value is ``False``.
        device (str): Specifying the device on which the optimization will be performed.
            Can be either 'cpu' or 'xpu' ('xpu' is not applicable for cpu only packages). The default value is 'cpu'.
        quantization_config (object): Defining the IPEX quantization recipe (Weight only quant or static quant).
            Default value is ``None``. Once used, meaning using IPEX quantizatization model for model.generate().
        qconfig_summary_file (str): Path to the IPEX static quantization config json file.
            Default value is ``None``. Work with quantization_config under static quantization use case.
            Need to do IPEX static quantization calibration and generate this file.
        low_precision_checkpoint (dict or tuple of dict): For weight only quantization with INT4 weights.
            If it's a dict, it should be the state_dict of checkpoint generated by GPTQ by default.
            If a tuple is provided, it should be `(checkpoint, quant_method)`,
            where `checkpoint` is the state_dict and `quant_method` is dict specifying the quantization
            method including GPTQ or AWQ, e,g, quant_method = {`quant_method`: `gptq`}.
        sample_inputs (Tuple tensors): sample inputs used for model quantization or torchscript.
            Default value is ``None``, and for well supported model, we provide this sample inputs automaticlly.
        deployment_mode (bool): Whether to apply the optimized model for deployment of model generation.
            It means there is no need to further apply optimization like torchscirpt. Default value is ``True``.
        cache_weight_for_large_batch (bool): Whether to cache the dedicated weight for large batch to speed up
            its inference (e.g., prefill phase) with extra memory usage. It is only valid for non-quantization cases
            where dtype = bfloat16 and weight-only quantization cases where lowp-mode=BF16/INT8. In other cases, an
            error will be raised. Default value is ``False``.

    Returns:
        Optimized model object for model.generate(), also workable with model.forward

    .. warning::
        Please invoke ``ipex.llm.optimize`` function AFTER invoking DeepSpeed in Tensor Parallel
        inference scenario.

    Examples:

        >>> # bfloat16 generation inference case.
        >>> model = ...
        >>> model.load_state_dict(torch.load(PATH))
        >>> model.eval()
        >>> optimized_model = ipex.llm.optimize(model, dtype=torch.bfloat16)
        >>> optimized_model.generate()

    """
    if isinstance(model, torch.jit.ScriptModule):
        return model
    if model.training or optimizer is not None:
        logger.warning(
            "fail to apply ipex.llm.optimize, this API supports inference for now, fallback to default path",
            _type=WarningType.NotSupported,
        )
        return model, optimizer

    validate_device_avaliable(device)

    try:
        well_supported_model = False
        if hasattr(model, "config") and hasattr(model.config, "architectures"):
            well_supported_model = model.config.architectures[0] in [
                "GPTJForCausalLM",
                "LlamaForCausalLM",
                "MllamaForConditionalGeneration",
                "GPTNeoXForCausalLM",
                "OPTForCausalLM",
                "FalconForCausalLM",
                "RWForCausalLM",
                "BloomForCausalLM",
                "CodeGenForCausalLM",
                "BaichuanForCausalLM",
                "ChatGLMModel",
                "GPTBigCodeForCausalLM",
                "T5ForConditionalGeneration",
                "MistralForCausalLM",
                "MixtralForCausalLM",
                "MptForCausalLM",
                "StableLmForCausalLM",
                "QWenLMHeadModel",
                "Qwen3MoeForCausalLM",
                "Qwen3ForCausalLM",
                "Qwen2ForCausalLM",
                "GitForCausalLM",
                "LlavaLlamaForCausalLM",
                "YuanForCausalLM",
                "PhiForCausalLM",
                "Phi3ForCausalLM",
                "Phi4MMForCausalLM",
                "WhisperForConditionalGeneration",
                "Maira2ForConditionalGeneration",
                "JambaForCausalLM",
                "DeepseekV2ForCausalLM",
                "DeepseekV3ForCausalLM",
            ]
        if well_supported_model:
            check_transformers_for_llm_support()
            if model.config.architectures[
                0
            ] == "Maira2ForConditionalGeneration" and hasattr(model, "language_model"):
                model.language_model = optimize(
                    model.language_model,
                    optimizer=optimizer,
                    dtype=dtype,
                    inplace=inplace,
                    device=device,
                    quantization_config=quantization_config,
                    qconfig_summary_file=qconfig_summary_file,
                    low_precision_checkpoint=low_precision_checkpoint,
                    sample_inputs=sample_inputs,
                    deployment_mode=False,
                    cache_weight_for_large_batch=cache_weight_for_large_batch,
                )
        else:
            if quantization_config is not None:
                logger.warning(
                    "ipex.llm.optimize supports quantizations on Llama, MLlama, GPT-J, GPT-Neox, Falcon, OPT, Bloom, CodeGen,"
                    + " Baichuan, ChatGLM, GPTBigCode, T5, Mistral, Mixtral, MPT, StableLM, QWen, Git, Llava, Yuan,"
                    + " Phi, Whisper, Maira2, Jamba and Deepseek, fallback to origin model"
                )
                return model

            if dtype is torch.float:
                _model = ipex.optimize(
                    model.eval(),
                    dtype=dtype,
                    inplace=inplace,
                    auto_kernel_selection=(
                        True
                        if ipex.get_fp32_math_mode() == ipex.FP32MathMode.BF32
                        else False
                    ),
                )
            elif dtype in [torch.bfloat16, torch.half]:
                _model = ipex.optimize(model.eval(), dtype=dtype, inplace=inplace)

            return _model

        if not inplace:
            _model = copy.deepcopy(model)
        else:
            _model = model

        # profiling mode is disabled in ChatGLM (https://huggingface.co/THUDM/chatglm3-6b/blob/main/modeling_chatglm.py#L33-L34)
        # Enable profiling mode to apply jit optimizations
        if model.config.architectures[0] == "ChatGLMModel":
            torch._C._jit_set_profiling_mode(True)
            torch._C._jit_set_profiling_executor(True)
            torch._C._jit_override_can_fuse_on_cpu(False)
            torch._C._jit_override_can_fuse_on_gpu(False)

        is_quantization = False
        is_woq = False
        if quantization_config is not None:
            is_quantization = True
            if _is_woq_qconfig(quantization_config):
                is_woq = True

        if cache_weight_for_large_batch:
            assert (
                is_woq or dtype == torch.bfloat16
            ), "cache_weight_for_large_batch is only valid for WOQ or BF16 cases"
            if is_woq:
                quantization_config = _woq_enable_weight_cache_for_large_batch(
                    quantization_config
                )

        # Load low precision checkpoint (generated by GPTQ, etc.) for WOQ before any conversion
        use_low_precision_checkpoint = False
        if device == "cpu" and is_woq and low_precision_checkpoint is not None:
            state_dict, quant_config = None, None
            if isinstance(low_precision_checkpoint, tuple):
                assert (
                    len(low_precision_checkpoint) == 2
                    and isinstance(low_precision_checkpoint[0], dict)
                    and isinstance(low_precision_checkpoint[1], dict)
                ), "Invalid low_precision_checkpoint"
                state_dict = low_precision_checkpoint[0]
                quant_config = low_precision_checkpoint[1]
            else:
                assert isinstance(
                    low_precision_checkpoint, dict
                ), "Invalid low_precision_checkpoint argument"
                state_dict = low_precision_checkpoint
                logger.warning(
                    "ipex.llm.optimize is loading low_precision_checkpoint state_dict"
                    " without quaquantization_method specified, using `gptq` by default."
                )
                quant_config = {
                    "quant_method": "gptq",
                    "group_size": None,
                }
            logger.debug(
                "ipex.llm.optimize is converting with low_precision_checkpoint"
            )
            _model = _convert_woq_with_low_precision_checkpoint(
                _model, quantization_config, state_dict, quant_config
            )
            use_low_precision_checkpoint = True

        # model reference conversion
        logger.debug("ipex.llm.optimize is converting model to reference model")
        _model = model_convert_reference(_model)

        # model quantization if needed
        if is_quantization:
            if not is_woq:  # static quantization
                if model.config.architectures[0] in ["MllamaForConditionalGeneration"]:
                    logger.warning(
                        "ipex.llm.optimize will skip static quantizations on MLlama ..."
                    )
                    return model

                deployment_mode = False

                if qconfig_summary_file is not None:
                    # static quantization is cpu only, thus doing model lowering firstly
                    _model = model_convert_lowering(
                        _model,
                        device,
                        dtype,
                        sample_inputs,
                        deployment_mode,
                        is_quantization,
                        woq=is_woq,
                    )
                    _model = ipex_quantization_flow(
                        _model,
                        dtype,
                        sample_inputs,
                        quantization_config,
                        qconfig_summary_file,
                    )
                    sample_inputs = (
                        get_dummy_input(_model, return_dict=True)
                        if sample_inputs is None
                        else sample_inputs
                    )
                    with torch.no_grad(), torch.cpu.amp.autocast(
                        enabled=(
                            True if dtype in [torch.bfloat16, torch.half] else False
                        ),
                        dtype=dtype,
                    ):
                        trace_model = torch.jit.trace(
                            _model,
                            example_kwarg_inputs=sample_inputs,
                            strict=False,
                            check_trace=False,
                        )
                        trace_model = torch.jit.freeze(trace_model)
                        if _model.config.architectures[0] == "YuanForCausalLM":
                            sample_inputs.pop("past_key_values", None)
                            batch_size = (
                                _model.config.batch_size
                                if hasattr(_model.config, "batch_size")
                                else 1
                            )
                            sample_inputs["input_ids"] = sample_inputs[
                                "input_ids"
                            ].repeat(batch_size, 1)
                            sample_inputs["attention_mask"] = sample_inputs[
                                "attention_mask"
                            ].repeat(batch_size, 1)
                            sample_inputs["position_ids"] = sample_inputs[
                                "position_ids"
                            ].repeat(batch_size, 1)
                            trace_model_first = torch.jit.trace(
                                _model,
                                example_kwarg_inputs=sample_inputs,
                                strict=False,
                                check_trace=False,
                            )
                            trace_model_first = torch.jit.freeze(trace_model_first)
                            _model = _set_optimized_model_for_generation(
                                _model,
                                optimized_model=trace_model,
                                first_token_optimized_model=trace_model_first,
                            )
                        else:
                            _model = _set_optimized_model_for_generation(
                                _model, optimized_model=trace_model
                            )
                    return _model
                else:
                    print(
                        "ipex.llm.optimize is prepared for the calibration of the static quantization"
                    )

            elif use_low_precision_checkpoint:  # weight only quantization
                # for non-quantized layers
                from ..frontend import optimize as ipex_optimize

                _model = _model.to(dtype)
                _model = ipex_optimize(
                    _model,
                    dtype=dtype,
                    inplace=True,
                    auto_kernel_selection=(dtype == torch.float),
                )
            else:
                # Note that low precision checkpoint is already handled at the beginning.
                # If checkpoint is not provided, model is quantized here
                _model = ipex_quantization_flow(
                    _model,
                    dtype,
                    sample_inputs,
                    quantization_config,
                    None,
                )
                # for fused moe
                if _model.config.architectures[0] in [
                    "DeepseekV2ForCausalLM",
                    "DeepseekV3ForCausalLM",
                    "Qwen3MoeForCausalLM",
                ]:
                    for layer in _model.model.layers:
                        layer.qconfig = quantization_config.global_qconfig

        # model lowering conversion
        logger.debug("ipex.llm.optimize is lowering model")
        _model = model_convert_lowering(
            _model,
            device,
            dtype,
            sample_inputs,
            deployment_mode,
            is_quantization,
            is_woq,
            cache_weight_for_large_batch,
        )
        # do not register output hook when doing calibration in static int8
        if not (is_quantization and not is_woq and qconfig_summary_file is None):
            from .models.reference.models import output_hook

            _model.register_forward_hook(output_hook, with_kwargs=True)
        return _model

    except RuntimeError as e:
        # if the model is not in the list, fallback without throwing errors
        if not well_supported_model:
            logger.warning(
                f"fail to apply ipex.llm.optimize due to: {e}, fallback to the origin model",
                _type=WarningType.NotSupported,
            )
            return model
        else:
            raise e

    return model

    Apply optimizations at Python frontend to the given transformers model (nn.Module).
    This API focus on transformers models, especially for generation tasks inference.

    Well supported model family with full functionalities:
    Llama, MLlama, GPT-J, GPT-Neox, OPT, Falcon, Bloom, CodeGen, Baichuan, ChatGLM, GPTBigCode,
    T5, Mistral, MPT, Mixtral, StableLM, QWen, Git, Llava, Yuan, Phi, Qwen3, Whisper. Maira2, Jamba, DeepSeekV2.

    For the model that is not in the scope of supported model family above, will try to
    apply default ipex.optimize transparently to get benifits (not include quantizations,
    only works for dtypes of torch.bfloat16 and torch.half and torch.float).

    Args:
        model (torch.nn.Module): User model to apply optimizations.
        optimizer (torch.optim.Optimizer): User optimizer to apply optimizations
            on, such as SGD. The default value is ``None``, meaning inference case.
        dtype (torch.dtype): Now it works for ``torch.bfloat16``, ``torch.half`` and ``torch.float``.
            The default value is ``torch.float``. When working with quantization, it means the mixed dtype with quantization.
        inplace (bool): Whether to perform inplace optimization. Default value is ``False``.
        device (str): Specifying the device on which the optimization will be performed.
            Can be either 'cpu' or 'xpu' ('xpu' is not applicable for cpu only packages). The default value is 'cpu'.
        quantization_config (object): Defining the IPEX quantization recipe (Weight only quant or static quant).
            Default value is ``None``. Once used, meaning using IPEX quantizatization model for model.generate().
        qconfig_summary_file (str): Path to the IPEX static quantization config json file.
            Default value is ``None``. Work with quantization_config under static quantization use case.
            Need to do IPEX static quantization calibration and generate this file.
        low_precision_checkpoint (dict or tuple of dict): For weight only quantization with INT4 weights.
            If it's a dict, it should be the state_dict of checkpoint generated by GPTQ by default.
            If a tuple is provided, it should be `(checkpoint, quant_method)`,
            where `checkpoint` is the state_dict and `quant_method` is dict specifying the quantization
            method including GPTQ or AWQ, e,g, quant_method = {`quant_method`: `gptq`}.
        sample_inputs (Tuple tensors): sample inputs used for model quantization or torchscript.
            Default value is ``None``, and for well supported model, we provide this sample inputs automaticlly.
        deployment_mode (bool): Whether to apply the optimized model for deployment of model generation.
            It means there is no need to further apply optimization like torchscirpt. Default value is ``True``.
        cache_weight_for_large_batch (bool): Whether to cache the dedicated weight for large batch to speed up
            its inference (e.g., prefill phase) with extra memory usage. It is only valid for non-quantization cases
            where dtype = bfloat16 and weight-only quantization cases where lowp-mode=BF16/INT8. In other cases, an
            error will be raised. Default value is ``False``.

    Returns:
        Optimized model object for model.generate(), also workable with model.forward

    .. warning::
        Please invoke ``ipex.llm.optimize`` function AFTER invoking DeepSpeed in Tensor Parallel
        inference scenario.

    Examples:

        >>> # bfloat16 generation inference case.
        >>> model = ...
        >>> model.load_state_dict(torch.load(PATH))
        >>> model.eval()
        >>> optimized_model = ipex.llm.optimize(model, dtype=torch.bfloat16)
        >>> optimized_model.generate()

    
intel_extension_for_pytorch/transformers/optimize.py