9j7ddlZddlZddlZddlmZddlZddlmZmZmZmZddl m Z ddl m Z m Z mZejZe dddd ej d edzd efd Zd eej(zdedededef dZdZedej0edeededededej8diZdZeZdej@d efdZ!dee"d eefdZ#y)N) GenericAlias)devicedtypeTensortypes) exposed_in)_resolve_opaque_type_infois_opaque_reference_typeis_opaque_typez torch.library)op_nameprototype_functionr returnc  d}|jdtj|fddtffd dtt t tzdfdtttjdftfffd d t t tzdttjtfffd g}t}d }tjjD]\}\}} t| sd | jtj j"k(r|s|j%dd}| j&tj j(ur d|d| j&\} } d} | t*vrt-| rt/| j0} n| t2j4j6k(r d|dnt9| dr_| j:turMt=| } d}| t*vrd| d}d|d| j&d|dt*j?d n8d|d| j&dt*j?dn t*| } | tAd|t |tur6||k7r tCd| jEdrJd|d | tGdd} n5||vr1| jEds d|d!d|d | tGdd} |jI|| jJtj j(ur|j%| d"|sd}| jJ%tM| jJtNtPtfrt| jJ}ntM| jJtt2jRfrd#| jJd#}ntM| jJt2jTrFt| jJ}d$}|jE|stAd%|d&|tG|d}n"d|d't | jJd(|j%| d"|d)|||k7r't||z }tG|d*kDr |d+jV\}} tY|}||d,d-j[|d.|Sd,d-j[|d.|S)/aJParses the schema of a given function with type hints. The schema is inferred from the function's type hints, and can be used to define a new operator. We make the following assumptions: * None of the outputs alias any of the inputs or each other. * | String type annotations "device, dtype, Tensor, types" without library specification are | assumed to be torch.*. Similarly, string type annotations "Optional, List, Sequence, Union" | without library specification are assumed to be typing.*. * | Only the args listed in ``mutates_args`` are being mutated. If ``mutates_args`` is "unknown", | it assumes that all inputs to the operator are being mutates. Callers (e.g. the custom ops API) are responsible for checking these assumptions. Args: prototype_function: The function from which to infer a schema for from its type annotations. op_name (Optional[str]): The name of the operator in the schema. If ``name`` is None, then the name is not included in the inferred schema. Note that the input schema to ``torch.library.Library.define`` requires a operator name. mutates_args ("unknown" | Iterable[str]): The arguments that are mutated in the function. Returns: The inferred schema. Example: >>> def foo_impl(x: torch.Tensor) -> torch.Tensor: >>> return x.sin() >>> >>> infer_schema(foo_impl, op_name="foo", mutates_args={}) foo(Tensor x) -> Tensor >>> >>> infer_schema(foo_impl, mutates_args={}) (Tensor x) -> Tensor unknownNc(td|dd)Nzinfer_schema(func): z Got func with signature )) ValueError)whatsigs [/media/conek/DATA/Code/OCR/venv/lib/python3.12/site-packages/torch/_library/infer_schema.pyerror_fnzinfer_schema..error_fnIs /v5NseSTUVVannotation_typecV t|S#t$rd|dYywxYw)NzUnsupported type annotation z. It is not a type.)eval Exception)rr pf_globals pf_localss rconvert_type_stringz)infer_schema..convert_type_stringLs< Y? ?  ..??RS  s ((tys.rcg}d}|D]#}|\}}|j|||z}%|r t|dfS|dfS)NFT)appendtuple)r!reschangedty ty_changedunstringify_types runstringify_typesz'infer_schema..unstringify_typesTs_ "B-b1NB JJrN z !G " :t# #: rr'ct|tr |dfStj|x}r.tj|\}}|rt ||dfS|dfS)NTF) isinstancestrtyping get_originget_argsr)r'originargs args_changedr r*s rr)z&infer_schema..unstringify_typebsl b# &r*D0 0((, ,V ,!26??23F!G D,#FD14775yrFz>We do not support positional-only args, varargs, or varkwargs.*Tz Parameter z must have a type annotation.z's type cannot be inferred from the schema as it is a ScriptObject. Please manually specify the schema using the `schema=` kwarg with the actual type of the ScriptObject. __origin__z z For example, z. z has unsupported type z]. We do not support Tuple inputs in schema. As a workaround, please try to use List instead. zThe valid types are: .. The valid types are: zschema_type is None for param zomutates_args must either be a sequence of the names of the arguments that are mutated or the string 'unknown'. rzTensor(az!)zM is in mutable_args but only Tensors or collections of Tensors can be mutated "ztorch.z dtype repr z must start with 'torch.'z' has an unsupported default value type z;. Please file an issue on GitHub so we can prioritize this.=rz in mutates_args were not found in the custom op's signature. mutates_args should contain the names of all args that the custom op mutates, or just the string 'unknown' if you don't know.(, z) -> ). __globals__inspect signaturer-r$typeobjectr.Anyboolset enumerate parametersitemssupported_paramkind Parameter KEYWORD_ONLYr# annotationemptySUPPORTED_PARAM_TYPESr r class_nametorch_C ScriptObjecthasattrr5 tuple_to_listkeysAssertionErrorr startswithlenadddefaultr,intfloatrrreturn_annotation parse_returnjoin)r mutates_argsr UNKNOWN_MUTATESparams seen_argssaw_kwarg_only_argidxnameparamr_ schema_type list_typeexample_type_str default_repr dtype_repr torch_dotmutates_args_not_seenr]retr rrrrr)r*s @@@@@@@r infer_schemarqs8T O#//JI   . /CWS  4<#%s* +  uVZZ_%t+ ,  T&\C/ E&**d:J4K FI'(<(<(>?[B]dEu% U V ::**77 7% c"%)"   w0066 6 z$'DE F.e.>.>? "7 7o.7HSS  EHH$9$99 'ZZ 6#..%7)/: #)  55)6yk'G$ &?!"5==/3 EMM5;;7 / $ !,,Y7(%j^3LM *#i.*:;  &MdSXS`S`NaMbcPQ MM[M4&,@ Aw[Bx& #L 1I = $ % ) ()*UV  ,C,A,ABq (( 3C!DIIf-.eC599 tyy !se ,,r base_typecpp_type list_baseoptional_base_listoptional_list_basecd|ftj|dfg}dttjzfd}|r |j fd||D|r1|j fd|tj|D|r |j fd||D|S)N?typctj|tj|ttj j|ftt |ffSN)r.SequenceListr collectionsabclist)rys rderived_seq_typesz(derived_types..derived_seq_typessF OOC KK  11C6 : v &   rc3,K|] }|df yw)z[]N.0seq_typrss r z derived_types..s! +2W "o & c3,K|] }|df yw)z?[]Nrrs rrz derived_types..s'  #& ' rc3NK|]}tj|dfyw)z[]?N)r.Optionalrs rrz derived_types..s. __W %(3'7 8 s"%)r.rr@ _SpecialFormextend)rrrsrtrurvresultrs ` r derived_typesrs H  #z^4KF  tf&9&99   6G 6R    -V__Y-GH    ,Y7   Mrc tddddftddddftddddftddddftddddft j ddddftd dddftd dddfg}tjjr4d d l m }|jtj t"|ddddfg}|D]}|j%t'|t)|S) NrTFSymIntr\rCr-Scalar ScalarTypeDevicer) GroupName)rr[r\rCr-rNumberrrrP distributed is_available"torch.distributed.distributed_c10drr#r.castr@rrdict)datarrlines rget_supported_param_typesrs 4u- heT* ud+ vtUD) eUE5) xue4  eUE2 5%/ LD %%'@ V[[y15%NO F, mT*+, <rrzTensor[]rr\rCrc|y|tjjur|dtj|}|t urD|t vr3t|rt|jS|d|dt dt |Stj|}|D]*}|t vs t|r|d|dt d,ddjfd|D}t|d k(rd |zd z}d |zd zS) Nz()z7No return type annotation was provided. Please add one.zReturn has unsupported type r7r6c^tj|x}r|St|jSr{)SUPPORTED_RETURN_TYPESgetr rO)argr's r_return_type_strz&parse_return.._return_type_strCs.'++C0 02 0I(-888rr<c3.K|] }|ywr{r)rrrs rrzparse_return..Hs@C*3/@sr r;r) r>rJrMr.r/r$rr r rOr0r_rX)rLrr1r2r output_tyrs @rr^r^(s'W&&,,,JK   z *F U 3 3' 30<GGG .zl;((>'?qB  &j11 ??: &D , ,5Mc5R .zl;((>'?qB 9  @4@@I 4yA~)Oc) ?S  rrgc|jtjjtjjfvSr{)rIr>rJPOSITIONAL_OR_KEYWORDrK)rgs rrHrHSs6 :://&& r tuple_typec.t|dd}|tjus|tus|dk(s|tSt |dk(r t|dSt |dk(r|dt ur t|dSttjt|S)z} Convert `tuple_type` into a list type with the same type arguments. Assumes that `tuple_type` is typing.Tuple type. __args__Nrr r)getattrr.Tupler$rrXEllipsisUnion)r type_argss rrTrTZs J5I fll"   ?   Y1 IaL!! Y1 1!9IaL!!FLLy!1233r)$r~r>r.rrrPrrrtorch.utils._exposed_inr opaque_objectr r r _TestTensorCallabler-rqr@rrCrrr}rr[r\rrr^rNrJrHr$rTrrrrsA  ...ll  O  C-C- 4Z C-  C-C-L$f)))$$$ $  $N2 H KKL* 7& LL(%!P237,,4d5k4d4j4r