
    ȅi2                         S SK r S SKJr  S SKJrJr  S SKJr  S SKrS SKJ	r	  SSK
7  SSKJrJr  SSK7  SSK7  SSK7  SSK7  SSK7  SSK7  SSK7  SSK7  SSK7  \" S	\\-  5      r/ S
QrS r " S S\5      rg)    N)Callable)OptionalUnion)TypeAliasType)Tensor   )*)fuse_modulesfuse_modules_qatObserverOrFakeQuantize)DeQuantStubFakeQuantizeFakeQuantizeBaseFixedQParamsFakeQuantizeFixedQParamsObserverFusedMovingAvgObsFakeQuantizeHistogramObserverMatchAllNodeMinMaxObserverMovingAverageMinMaxObserver%MovingAveragePerChannelMinMaxObserverNoopObserverObserverBaser   PatternPerChannelMinMaxObserverPlaceholderObserverQConfig
QConfigAnyQConfigDynamicQConfigMapping	QuantStub	QuantTypeQuantWrapperRecordingObserverReuseInputObserverUniformQuantizationObserverBaseadd_quant_dequantconvertconvert_dynamic_jitconvert_jit'default_affine_fixed_qparams_fake_quant%default_affine_fixed_qparams_observerdefault_debug_observerdefault_dynamic_fake_quantdefault_dynamic_quant_observerdefault_embedding_fake_quant!default_embedding_fake_quant_4bitdefault_eval_fndefault_fake_quant+default_fixed_qparams_range_0to1_fake_quant)default_fixed_qparams_range_0to1_observer.default_fixed_qparams_range_neg1to1_fake_quant,default_fixed_qparams_range_neg1to1_observerdefault_float_qparams_observer#default_float_qparams_observer_4bitdefault_fused_act_fake_quant'default_fused_per_channel_wt_fake_quantdefault_fused_wt_fake_quantdefault_histogram_fake_quantdefault_histogram_observerdefault_observer%default_per_channel_weight_fake_quant#default_per_channel_weight_observerdefault_placeholder_observerdefault_reuse_input_observer*default_symmetric_fixed_qparams_fake_quant(default_symmetric_fixed_qparams_observerdefault_weight_fake_quantdefault_weight_observerdisable_fake_quantdisable_observerenable_fake_quantenable_observerfuse_conv_bnfuse_conv_bn_jitfuse_conv_bn_relufuse_convtranspose_bnfuse_linear_bnr
   r   4fused_per_channel_wt_fake_quant_range_neg_127_to_127(fused_wt_fake_quant_range_neg_127_to_127get_combined_dict&get_default_compare_output_module_listget_default_custom_config_dict)get_default_dynamic_quant_module_mappings0get_default_dynamic_sparse_quant_module_mappings0get_default_float_to_quantized_operator_mappingsget_default_qat_module_mappingsget_default_qat_qconfigget_default_qat_qconfig_dictget_default_qat_qconfig_mappingget_default_qconfigget_default_qconfig_dictget_default_qconfig_mapping$get_default_qconfig_propagation_list(get_default_static_quant_module_mappings2get_default_static_quant_reference_module_mappings/get_default_static_sparse_quant_module_mappingsget_dynamic_quant_module_class!get_embedding_qat_module_mappings*get_embedding_static_quant_module_mappingsget_fuser_methodget_fuser_method_newget_observer_state_dictget_quantized_operatorget_static_quant_module_classload_observer_state_dictno_observer_set0per_channel_weight_observer_range_neg_127_to_127prepareprepare_dynamic_jitprepare_jitprepare_qatpropagate_qconfig_qconfig_equalsquantizequantize_dynamicquantize_dynamic_jitquantize_jitquantize_qatscript_qconfigscript_qconfig_dictswap_module$weight_observer_range_neg_127_to_127AffineQuantizedObserverBaseGranularityMappingTypePerAxisPerBlockPerGroupPerRow	PerTensorPerTokenTorchAODTypeZeroPointDomainget_block_sizec                 *    U H  u  p#U " U5        M     g)zDefine the default evaluation function.

Default evaluation function takes a torch.utils.data.Dataset or a list of
input Tensors and run the model on the dataset
N )model
calib_datadata_targets       X/home/james-whalen/.local/lib/python3.13/site-packages/torch/ao/quantization/__init__.pyr2   r2      s     $d $    c                      ^  \ rS rSrSr    SS\R                  S\\   S\	\\   /\
\\4   4   S\S-  S\S-  S	\R                  S-  S
\S-  4U 4S jjjrS\S\4S jrS rSrU =r$ )_DerivedObserverOrFakeQuantize   zmThis observer is used to describe an observer whose quantization parameters
are derived from other observers
Ndtype
obs_or_fqsderive_qparams_fn	quant_min	quant_maxqschemech_axisc                    > [         T	U ]  U5        X l        X0l        X@l        XPl        X`l        Xpl        SSKJ	n  U" U R                  5      (       a  U R                  c  [        S5      eg g )Nr   )is_per_channelz6Must provide a valid ch_axis if qscheme is per channel)super__init__r   r   r   r   r   r   utilsr   AssertionError)
selfr   r   r   r   r   r   r   r   	__class__s
            r   r   '_DerivedObserverOrFakeQuantize.__init__   sg     	$!2"")$,,''||#$L  $ (r   xreturnc                     U$ Nr   )r   r   s     r   forward&_DerivedObserverOrFakeQuantize.forward   s    r   c                 8    U R                  U R                  5      $ r   )r   r   )r   s    r   calculate_qparams0_DerivedObserverOrFakeQuantize.calculate_qparams   s    %%doo66r   )r   r   r   r   r   r   )NNNN)__name__
__module____qualname____firstlineno____doc__torchr   listr   r   tupler   intr   r   r   r   __static_attributes____classcell__)r   s   @r   r   r      s     !% $(,"{{ /0 $()*E&&.,AA
	 : : % t 8 F 7 7r   r   )syscollections.abcr   typingr   r   typing_extensionsr   r   r   fake_quantizer
   r   fuser_method_mappingsobserverqconfigqconfig_mapping
quant_typequantization_mappingsru   rx   stubsr   r   r   __all__r2   r   r   r   r   <module>r      sl     $ " +    8 $     $    'l-== 
CL%7\ %7r   