
    sh                    j    S SK Jr  S SKJr  SSKJrJr  \(       a
  S SKrS SKJ	r	  S/r
 " S S\5      rg)	    )annotations)TYPE_CHECKING   )QuantizationAnnotation	QuantizerN)NodeComposableQuantizerc                  l   ^  \ rS rSrSrS	U 4S jjr      S
S jrSS jr    SS jrSS jr	Sr
U =r$ )r	      a4  
ComposableQuantizer allows users to combine more than one quantizer into a single quantizer.
This allows users to quantize a model with multiple quantizers. E.g., embedding quantization
maybe supported by one quantizer while linear layers and other ops might be supported by another
quantizer.

ComposableQuantizer is initialized with a list of `Quantizer` instances.
The order of the composition matters since that is the order in which the quantizers will be
applies.
Example:
```
embedding_quantizer = EmbeddingQuantizer()
linear_quantizer = MyLinearQuantizer()
xnnpack_quantizer = XNNPackQuantizer() # to handle ops not quantized by previous two quantizers
composed_quantizer = ComposableQuantizer([embedding_quantizer, linear_quantizer, xnnpack_quantizer])
prepared_m = prepare_pt2e(model, composed_quantizer)
```
c                <   > [         TU ]  5         Xl        0 U l        g N)super__init__
quantizers_graph_annotations)selfr   	__class__s     ڈ/Users/tiagomarins/Projetos/claudeai/copy_bank/venv/lib/python3.13/site-packages/torch/ao/quantization/quantizer/composable_quantizer.pyr   ComposableQuantizer.__init__%   s    $FH    c                   UR                   R                   H  nSUR                  ;   a  X0R                  ;   aW  [	        U R                  U   5      [	        UR                  S   5      :w  a%  [        SUR                  R                   SU 35      eUR                  S   U R                  U'   M  X0R                  ;   d  M  [        SUR                  R                   SU 35      e   g )Nquantization_annotationz
Quantizer z! has changed annotations on node z! has removed annotations on node )graphnodesmetar   idRuntimeErrorr   __name__)r   gm	quantizerns       r    _record_and_validate_annotations4ComposableQuantizer._record_and_validate_annotations*   s     A(AFF2 ///t..q12!&&!:;<= '$Y%8%8%A%A$BBcdecfg  238Q1RD++A.///&$Y%8%8%A%A$BBcdecfg   r   c                p    U R                    H%  nUR                  U5        U R                  X5        M'     U$ )z!just handling global spec for now)r   annotater"   r   modelr    s      r   r%   ComposableQuantizer.annotate@   s3    Iu%11%C ) r   c                N    U R                    H  nUR                  U5      nM     U$ r   )r   transform_for_annotationr&   s      r   r*   ,ComposableQuantizer.transform_for_annotationG   s'     I66u=E )r   c                    g r    )r   r'   s     r   validateComposableQuantizer.validateN   s    r   )r   r   )r   zlist[Quantizer])r   torch.fx.GraphModuler    r   returnNone)r'   r0   r1   r0   )r'   r0   r1   r2   )r   
__module____qualname____firstlineno____doc__r   r"   r%   r*   r.   __static_attributes____classcell__)r   s   @r   r	   r	      sO    &I
&3<	,)	 r   )
__future__r   typingr   r    r   r   torchtorch.fxr   __all__r	   r-   r   r   <module>r>      s1    "   8  
>) >r   