Vh$ddlZddlmZddlmZgdZGddej j ZGddej j ZGd d ej j Z y) N)Tensor)ops)FloatFunctionalFXFloatFunctional QFunctionalceZdZdZdfd ZdZ dededefdZ dededefdZ dededefd Z dededefd Z dde ed e defd Z dededefd Z dededefdZxZS)raState collector class for float operations. The instance of this class can be used instead of the ``torch.`` prefix for some operations. See example usage below. .. note:: This class does not provide a ``forward`` hook. Instead, you must use one of the underlying functions (e.g. ``add``). Examples:: >>> f_add = FloatFunctional() >>> a = torch.tensor(3.0) >>> b = torch.tensor(4.0) >>> f_add.add(a, b) # Equivalent to ``torch.add(a, b)`` Valid operation names: - add - cat - mul - add_relu - add_scalar - mul_scalar returncht|tjj |_yN)super__init__torchnnIdentityactivation_post_processself __class__s `/home/dcms/DCMS/lib/python3.12/site-packages/torch/ao/nn/quantized/modules/functional_modules.pyr zFloatFunctional.__init__&s# ',xx'8'8':$ctdNzYFloatFunctional is not intended to use the 'forward'. Please use the underlying operation RuntimeErrorrxs rforwardzFloatFunctional.forward* ?  rrycTtj||}|j|}|Sr )raddrrrrrs rr!zFloatFunctional.add2& IIaO  ( ( +rc2tj||}|Sr rr!r"s r add_scalarzFloatFunctional.add_scalar9 IIaOrcTtj||}|j|}|Sr )rmulrr"s rr*zFloatFunctional.mulAr$rc2tj||}|Sr rr*r"s r mul_scalarzFloatFunctional.mul_scalarHr(rdimcVtj||}|j|}|SN)r.)rcatrrrr.r#s rr1zFloatFunctional.catPs' IIaS !  ( ( +rctj||}tjjj |}|j |}|Sr )rr!r functionalrelurr"s radd_reluzFloatFunctional.add_reluWs@ IIaO HH   $ $Q '  ( ( +rcTtj||}|j|}|Sr )rmatmulrr"s rr8zFloatFunctional.matmul_s' LLA   ( ( +rr Nr)__name__ __module__ __qualname____doc__r rrr!floatr'r*r-listintr1r6r8 __classcell__rs@rrr s4; AV6 @Fu AV6 @Fu 1T&\F <&V D6frrceZdZdZdZ dededefdZ dededefdZ dededefdZ dededefd Z dde ed e defd Z dededefd Z dededefd Zy)ra#module to replace FloatFunctional module before FX graph mode quantization, since activation_post_process will be inserted in top level module directly Valid operation names: - add - cat - mul - add_relu - add_scalar - mul_scalar ctdrrrs rrzFXFloatFunctional.forwardrrrrrr c2tj||}|Sr r&r"s rr!zFXFloatFunctional.addz IIaOrc2tj||}|Sr r&r"s rr'zFXFloatFunctional.add_scalarrGrc2tj||}|Sr r,r"s rr*zFXFloatFunctional.mulrGrc2tj||}|Sr r,r"s rr-zFXFloatFunctional.mul_scalarrGrr.c4tj||}|Sr0)rr1r2s rr1zFXFloatFunctional.cats IIaS !rctj||}tjjj |}|Sr )rr!rr4r5r"s rr6zFXFloatFunctional.add_relus0 IIaO HH   $ $Q 'rc2tj||}|Sr )rr8r"s rr8zFXFloatFunctional.matmuls LLA rNr:)r;r<r=r>rrr!r?r'r*r-r@rAr1r6r8rrrres  AV6@FuAV6@Fu1T&\F<&V D6frrceZdZdZdfd ZfdZfdZdZdZdZ d e d e de fd Z d e d e de fd Z d e d e de fd Z d e d e de fdZ dd ee dede fdZ d e d e de fdZ d e d e de fdZeddZxZS)ra Wrapper class for quantized operations. The instance of this class can be used instead of the ``torch.ops.quantized`` prefix. See example usage below. .. note:: This class does not provide a ``forward`` hook. Instead, you must use one of the underlying functions (e.g. ``add``). Examples:: >>> q_add = QFunctional() >>> # xdoctest: +SKIP >>> a = torch.quantize_per_tensor(torch.tensor(3.0), 1.0, 0, torch.qint32) >>> b = torch.quantize_per_tensor(torch.tensor(4.0), 1.0, 0, torch.qint32) >>> q_add.add(a, b) # Equivalent to ``torch.ops.quantized.add(a, b, 1.0, 0)`` Valid operation names: - add - cat - mul - add_relu - add_scalar - mul_scalar r ct|d|_d|_tj j |_y)Ng?r)r r scale zero_pointrrrrrs rr zQFunctional.__init__s1  ',xx'8'8':$rct||||tj|j||dz<tj|j ||dz<y)NrQrR)r _save_to_state_dictrtensorrQrR)r destinationprefix keep_varsrs rrTzQFunctional._save_to_state_dictsL #KC(- TZZ(@ FW$%-2\\$//-J F\)*rc t|j|dz|_t|j|dz|_t ||||d|||y)NrQrRF)r?poprQrArRr _load_from_state_dict) r state_dictrWlocal_metadatastrict missing_keysunexpected_keys error_msgsrs rr[z!QFunctional._load_from_state_dicts[:>>&7*:;< jnnVl-BCD %        rcy)NrrNrs r _get_namezQFunctional._get_namesrc:d|jd|jS)Nzscale=z , zero_point=rQrRrcs r extra_reprzQFunctional.extra_reprs |=0ABBrctd)NzTFunctional is not intended to use the 'forward'. Please use the underlying operationrrs rrzQFunctional.forwardrrrrctjj|||j|j}|j |}|SNrf)r quantizedr!rQrRrr"s rr!zQFunctional.add< MM  a$**  Q  ( ( +rcFtjj||}|Sr )rrkr'r"s rr'zQFunctional.add_scalar MM $ $Q *rctjj|||j|j}|j |}|Srj)rrkr*rQrRrr"s rr*zQFunctional.mulrlrcFtjj||}|Sr )rrkr-r"s rr-zQFunctional.mul_scalarrnrr.ctjj||j|j|}|j |}|S)N)rQrRr.)rrkr1rQrRrr2s rr1zQFunctional.cat s= MM  atzzdooSV  W  ( ( +rctjj|||j|j}|j |}|Srj)rrkr6rQrRrr"s rr6zQFunctional.add_relus< MM " "1atzzdoo " V  ( ( +rcttjj|||j|j}|Srj)rrkr8rQrRr"s rr8zQFunctional.matmuls. MM ATZZDOO Trct|tk(sJd|jj\}}t }t ||_t||_|S)Nz=QFunctional.from_float expects an instance of FloatFunctional) typerrcalculate_qparamsrr?rQrArR)clsmoduse_precomputed_fake_quantrQrRnew_mods r from_floatzQFunctional.from_float!s^ I ( K J K (77IIKz-e   _rr9r:)F)r;r<r=r>r rTr[rdrgrrr!r?r'r*r-r@rAr1r6r8 classmethodr{rBrCs@rrrs6; K  ,C ?V6 NFu OV6 NFu ?T&\F D&V R6f rr) rr torch._opsr__all__rModulerrrrNrrrsZ  BWehhooWt<<~F%((//Fr