o
    pi                     @   sD   d dl Z ddlmZ ddlmZ G dd deZG dd deZdS )	    N   )BaseObserver)ObserverFactoryc                       s*   e Zd ZdZd fdd	Zdd Z  ZS )AbsmaxObservera  
    It collects maximum absolute values of target tensor.

    Args:
        bit_length(int, optional): Number of bits to represent an quantized integer in binary.
        dtype(str, optional): The data type of input tensor.
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.

    Examples:
        .. code-block:: python

            >>> from paddle.quantization import QuantConfig
            >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
            >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.99)
            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
       c                    s   t  j|d d S )N)
quant_bits)super__init__)selfr   	__class__ l/home/app/PaddleOCR-VL/.venv_paddleocr/lib/python3.10/site-packages/paddle/quantization/observers/abs_max.pyr	   )   s   zAbsmaxObserver.__init__c                 C   s   t S N)AbsmaxObserverLayerr
   r   r   r   
_get_class,      zAbsmaxObserver._get_classr   )__name__
__module____qualname____doc__r	   r   __classcell__r   r   r   r   r      s    r   c                       s^   e Zd ZdZdZd fdd	Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Z  ZS )r   z'
    Per-tensor abs max quantizer.
    gHz>r   c                    s4   t    || _ttj| _d | _d | _	d | _
d S r   )r   r	   _quant_bitspaddleZ	to_tensorr   INIT_ABS_MAXabs_max_val_max_scale_zero_point)r
   layerr   r   r   r   r	   7   s   

zAbsmaxObserverLayer.__init__c                 C   s   |  |\| _| _|S r   )cal_min_max_minr   )r
   inputr   r   r   forward?   s   zAbsmaxObserverLayer.forwardc                 C   s8   t t |}| jd urt || j|j}d|fS )Nr   )r   maxabsr   maximumcastZdtype)r
   Zinputsr   r   r   r   r"   C   s   
zAbsmaxObserverLayer.cal_min_maxc                 C      | j S r   )r   r   r   r   r   
bit_lengthK   s   zAbsmaxObserverLayer.bit_lengthc                 C   s   dS )Nr   r   r   r   r   
quant_axisN   r   zAbsmaxObserverLayer.quant_axisc                 C   s$   | j du r	| j| _ t| j | _dS )z$Compute thresholds for MAX function.N)r   r   r   Z
zeros_liker    r   r   r   r   cal_thresholdsQ   s   
z"AbsmaxObserverLayer.cal_thresholdsc                 C   s   | j du r	|   | j S )zReturn output scales.N)r   r.   r   r   r   r   scalesW   s   
zAbsmaxObserverLayer.scalesc                 C   r*   )zReturn output zero points.)r    r   r   r   r   zero_points]   s   zAbsmaxObserverLayer.zero_pointsr   )r   r   r   r   r   r	   r%   r"   r+   r-   r.   r/   r0   r   r   r   r   r   r   0   s    r   )r   Zbase_observerr   factoryr   r   r   r   r   r   r   <module>   s
   