a hi@sUddlZddlZddlZddlZddlZddlZddlZddlZddlm Z ddl Zddl m Z ddl m Z ddl mZddlmZmZmZmZmZmZddlmZmZgdZdd d Zhd Zd d Zd:eejeejdddZ iZ!e"edfe#d<GdddZ$ej%GdddZ&eedddZ'eddddZ(eddddZ)ee*eefdd d!Z+edd"d#d$Z,ej-e.d%d&d'Z/eejdd(d)d*Z0ej1edd+d,d-Z2d.d/Z3d0d1Z4d;d3d4Z5d5d6Z6dtorchatZprimsZprimZatenZpytorchcCstdtdS)Nzwtorch._custom_op is deprecated and will be removed in PyTorch 2.6, please use the equivalent torch.library API instead.)warningswarnDeprecationWarningrrQ/var/www/html/assistant/venv/lib/python3.9/site-packages/torch/_custom_op/impl.pywarn_deprecated5sr)qualname manual_schemareturncstfdd}|S)L This API is deprecated, please use torch.library.custom_op instead c s,t|stdt|t\}}t||j|krXtdd|d|jddurlt|ddn}||}t |}t |durt ||t |d}||t||j}t|||||d d }|j|_|j|_|j|_t ||jd tt|tj|ttt||S) NzDcustom_op(...)(func): Expected `func` to be a Python function, got: zcustom_op(qualname='z-', ...)(func): expected `func` to have name 'z ' but got 'zX'. Please either change the name of `func` or the qualname that is passed to `custom_op`r)Z mutates_argsFRAGMENTT_private_accessAutograd)inspect isfunction ValueErrortypeparse_qualnamevalidate_namespace__name__rrparsevalidate_schema validate_function_matches_schemalibraryLibrarydefinefind_ophandle_or_thrownamer __module____doc__impl_opnamer weakrefproxyr_C#_dispatch_set_report_error_callback functoolspartialreport_error_callback) funcnsr1schema schema_strfunction_schemalibophandleresultrrrrinnerEsR            zcustom_op..inner)r)rrrFrrErr =s1r rglobal_registrycseZdZdZddfdd ZddZd,d d Zd d Zd dZddZ ddZ ddZ d-e j ee jefe jdddZddZe jdddZd.e jdddZdd Zd!d"Zd#d$Zd%d&Zd/d'd(Zd0d*d+ZZS)1rrFr csntt|std|d|}||_||_||_||_||_||_ d|_ i|_ d|_ |t |j <dS)Nz|The CustomOp constructor is private and we do not guarantee BC for it. Please use custom_op(...) to create a CustomOp object::F)super__init__r RuntimeError_schema_cpp_ns_lib _ophandler5 _qualnamer)_impls'_registered_autograd_kernel_indirectionrG)selfrBcpp_nsr? operator_namerCr!r1 __class__rrrJs" zCustomOp.__init__cCs0|jr J|j|jtt|dd|_dS)Nr"T)rRrNr4r5r r6r7rSrrr%_register_autograd_kernel_indirections  z.CustomOp._register_autograd_kernel_indirectionc Cs||rJ|j|}|dus J|j}td|d|jd|d|d tt|}|j d|j }t |||j|<dS)NzAttempting to register a z impl for operator z that already has a z impl registered from Python at z. This is not supported.:) _has_implrQlocationrKrPr# getframeinfosys _getframefilenamelinenoFuncAndLocation)rSkindr= stacklevelZfunc_and_locationr]framerrr_register_impls   zCustomOp._register_implcCs |j|SNrQrSrdrrr _get_implszCustomOp._get_implcCs ||jvSrhrirjrrrr\szCustomOp._has_implcCs6|`ttj|j}t||jr*t||jt|j =dSrh) rNgetattrropsrMhasattrr5delattrrGrP)rS opnamespacerrr_destroys   zCustomOp._destroycCsd|jdS)Nz)rPrXrrr__repr__szCustomOp.__repr__cOstj|jg|Ri|}|Srh)r8Z_dispatch_call_boxedrO)rSargskwargsrDrrr__call__szCustomOp.__call__) device_typesrcs6ttrgD] }t|qfdd}|S)T This API is deprecated, please use torch.library.custom_op instead csJtD]<}|j||dt|}tjj||q|S)Nre)set_check_doesnt_have_library_implrgSUPPORTED_DEVICE_TYPE_TO_KEYr-r4rNr5)f device_typeZ dispatch_key _stacklevelrvrSrrrFs   zCustomOp.impl..inner) isinstancestrvalidate_device_type)rSrvrr}rFrr~rr4s   z CustomOp.implcCs@||rdSt|}t|j|r.innerr)rSrFrrXr impl_factorys zCustomOp.impl_factorycsfdd}|S)rwcsZjdddjjtfdd}jj |dS)NabstractrxcsNfdd}tjj||i|WdS1s@0YdS)NcstdddS)Nz.inner..f_with_ctx..error_on_ctx)r_libraryZ fake_implZset_ctx_getter)rsrtrr|r]rrr f_with_ctxs z9CustomOp.impl_abstract..inner..f_with_ctxMeta) $_check_doesnt_have_library_meta_implrgrkr]rPr:wrapsrNr4r5)r|rrrSrrrF s z%CustomOp.impl_abstract..innerrrSrrFrrr impl_abstractszCustomOp.impl_abstractc sfdd}jtjkr*|dj}js>|dt|dksNJtdd|D}|rl|dttj d ttj d ttj d ttj d ttj d tttj ddi}jD]0}|j|vrq|dt|d|jdqdS)Ncs td|djddS)NzCCannot use torch._custom_ops APIs to register backward formula for z. Got operator z with schema: )rKrP)detailr?rSrrerror*sz4CustomOp._check_can_register_backward..errorznon-functional operatorzoperator with no returnsrcss"|]}|jduo|jj VqdSrh) annotationZis_write).0rrrr :sz8CustomOp._check_can_register_backward..zoperator that returns viewsintSymIntboolfloatTensorz List[Tensor]zoperator with return not in z (got ))rLrdr Z functionalreturnslenanyrrrrrrrrr&listvalues)rSrZretsZis_non_mutating_viewZallowed_return_typesretrrr_check_can_register_backward)s4       z%CustomOp._check_can_register_backwardcCs^|jr dSt|jdr*td|jddD]*}t|j|r.td|jd|dq.dS)NCompositeImplicitAutogradz3impl_backward/impl_save_for_backward: the operator a3 already has an implementation for this device type via a pre-existing registration to DispatchKey::CompositeImplicitAutograd.CompositeImplicitAutograd operators do not need an autograd formula; instead, the operator will decompose into its constituents and those can have autograd formulas defined on them.)r"Z AutogradCPUZ AutogradCUDAz; already has an Autograd kernel registered to DispatchKey::z vi a pre-existing torch.library or TORCH_LIBRARY registration. Please either remove those registrations or don't use the torch._custom_ops APIs)rRr8%_dispatch_has_kernel_for_dispatch_keyrPrK)rSrrrr(_check_doesnt_have_library_autograd_implPs"  z1CustomOp._check_doesnt_have_library_autograd_implcCsr|drdSt|jdr.t|jds.dSt|jdrNtd|jdt|jdrntd|jddS)NrZCompositeExplicitAutogradrrz!impl_abstract(...): the operator a- already has an implementation for this device type via a pre-existing registration to DispatchKey::CompositeImplicitAutograd.CompositeImplicitAutograd operators do not need an abstract impl; instead, the operator will decompose into its constituents and those can have abstract impls defined on them.z already has an DispatchKey::Meta implementation via a pre-existing torch.library or TORCH_LIBRARY registration. Please either remove that registration or don't call impl_abstract.)r\r8rrPrKrXrrrrls$     z-CustomOp._check_doesnt_have_library_meta_implc CsX|dsJ|dsJt|j|j|t|j|dj|dj}|d|dS)Nbackwardsave_for_backwardautograd) r\r rL_output_differentiabilityget_oprPrkr=rg)rSZkernelrrr_register_autograd_kernels  z"CustomOp._register_autograd_kernelcsfdd}|S)zyRegister a function that tells us what to save for backward. Please see impl_backward for more details. csDjsjd|ddr@dS)Nrrxr)rrrRrYrgr\rrrrrrFs z.CustomOp.impl_save_for_backward..innerrrrrrimpl_save_for_backwards zCustomOp.impl_save_for_backwardNcsldurXfdd}tts$|D]}t|ts(|q(tjjtkrX|fdd}|S)rwNcstddS)Nzimpl_backward(output_differentiability): expected output_differentiability to be a list of bools with length equal to the number of outputs of this CustomOp got: rr)output_differentiabilityrryells z$CustomOp.impl_backward..yellcsJjsjd|d_drFdS)Nrrxr)rrrRrYrgrr\rrrrrSrrrFs z%CustomOp.impl_backward..inner)rrrrrLr)rSrrrdiffrFrrr impl_backwards    zCustomOp.impl_backward)rZ)rZ)rZ)rZ)NrZ)r)r2 __qualname__r3rJrYrgrkr\rqrrrutypingUnionrIterableCallabler4rzrrrrrrrr __classcell__rrrVrrs0     #', c@s eZdZUejed<eed<dS)rcr=r]N)r)r2rrr__annotations__rrrrrrcs  rc)rTrUcCs0|jdurdn|j}t|dt|j|S)NrH) overload_namer8Z_dispatch_find_schema_or_throwrr1)rTrUrrrrr0s r0)r>rcCs:d|vrtd|d|tvr6td|d|ddS)N.zcustom_op(..., ns="zC"): expected ns to not contain any . (and be a valid variable name)zcustom_op(..., ns='z'): 'z9' is a reserved namespace, please choose something else. )r% RESERVED_NS)r>rrrr(s r()r?rcCs:tjj|std||jjdur6td|dS)Nzcustom_op only supports functional operators (ops that do not mutate any inputs, do not return views of the inputs, and has at least one return). Got the following non-functional schema: zUcustom_op does not support arguments named 'self'. Please rename your argument. Got: )rrutilsZis_functional_schemar% argumentsZself_arg)r?rrrr+s r+)rrcCsR|dd}t|dkr(td|dd|dvrBtd||d|dfS) NrHr rZz$Expected there to be a namespace in z;, i.e. The operator name should look something like ns::foorzThe torch.custom_ops APIs do not handle overloads, i.e. operator names with '.' in them. Please name your operator something like ns::foo. Got: r)splitrr%)rnamesrrrr's    r')r}rcCs&|tvr"td|dtddS)NzCustomOp.impl(device_types=[z(, ...]): we only support device_type in r)r{r%keys)r}rrrrs r)paramrcCs|jtjjtjjfvSrh)rdr# ParameterPOSITIONAL_OR_KEYWORD KEYWORD_ONLY)rrrrsupported_paramsr)r?r=rcst|tddjDs0tdtddjDsVjtjj urdtdddjD}ddjD}fd d fd d fd d}||j j ||j j dS)Ncss|]\}}t|VqdSrh)rr_prrrr+z3validate_function_matches_schema..zcustom_op(..., manual_schema)(func): positional-only args, varargs, and kwargs are not supported. Please rewrite `func` to not have them. Got `func` with signature: css |]\}}|jtjjuVqdSrh)rr#remptyrrrrr3szcustom_op(..., manual_schema)(func): When passing in a manual schema, we expect `func` to have no type annotations to avoid ambiguity. Got `func` with signature: cSs&g|]\}}|jtjjkr||fqSr)rdr#rrrr1rrrr ?sz4validate_function_matches_schema..cSs&g|]\}}|jtjjkr||fqSr)rdr#rrrrrrrDscstdddS)Nzcustom_op(..., manual_schema)(func): When passing in a manual schema, we expect `func`'s signature to match `manual_schema` (aside from type annotations). func's signature: , manual_schema: r%rr?sigrrrJsz/validate_function_matches_schema..errorcstdddS)Nzycustom_op(..., manual_schema)(func): neither func nor manual_schema should have default arguments. Got func's signature: rrrrrrerror_default_argsRsz.error_default_argscs`t|t|krt||D]:\\}}}||jkr<|jtjjusT|jdur q dSrh)rzipr1defaultr#rr)Zsig_argsZ schema_argsr1rarg)rrrrcompareZs z1validate_function_matches_schema..compare) r# signatureall parametersitemsr%rreturn_annotation SignaturerrZflat_positionalZflat_kwarg_only)r?r= positionalZ kwargonlyrr)rrr?rrr,&s:   r,)r rrc Cst|dkrt|d|dkr,t|d|dvr\|}t|d|d|d|d t|d |d dS) NZ Undefineda: There were no Tensor inputs to this operator (e.g. you passed an empty list of Tensors). If your operator is a factory function (that is, it takes no Tensors and constructs a new one), then please use CustomOp.impl_factory to register an implementation for itrz: when running with device='Meta' tensors: there is no abstract impl registered for this CustomOp. Please register one via CustomOp.impl_abstract to get this CustomOp to work with Meta tensors)rrz: when running with device='z' tensors: there is no zW impl registered for this CustomOp. Please register one via CustomOp.impl(device_type='z')z%: No implementation for dispatch key z. It is likely that we have not added this functionality yet, please either open an issue or if you're feeling adventurous, use the low-level torch.library API)NotImplementedErrorlower)r rZdevicerrrr<gs( r<cCs\|j}tj|d}|dd}t|j}|dd}t |}t |||||ddS)NrrHTr ) namespacerr-r.r1rrrLrr*r)opr>rBr1r@r?rrrcustom_op_from_existings  rcsffdd}t\}}ttj|s*|ttj|}t||sF|t||}t|ds`||jS)NcstdddS)NzCould not find the operator z~. Please make sure you have already registered the operator and (if registered from C++) loaded it via torch.ops.load_library.rrrrrerror_not_founds zget_op..error_not_foundr)r'rnrrmrlr)rrr>r1rppacketrrrrs       rFcCs8|tvrt|S|s$td|dt|}t|}|S)NzCould not find custom op "z5". Did you register it via the torch._custom_ops API?)rGrKrr)rZalso_check_torch_libraryoverloadrDrrr_find_custom_ops rcCsF|tjjjvrdStjjj|}|dur,dS|ds:dS|djS)Nr)rZ _custom_opr4rGr\rkr=)rr rrrget_abstract_impls rTc Cs|d\}}||}t|}t||rr1r@rArrBrCrDrrr_custom_op_with_schemas    r)N)F)T)8 dataclassesr:r#r_rrr6rZtorch._Cr8Ztorch._library.infer_schemaZ torch.libraryr-rrZtorchgen.modelrrrrrr rr r __all__r{rrrOptionalrr rGdictrr dataclassrcr0r(r+tupler'rrrrr,Anyr<rrrrrrrrrs\          BY   A