o vZh@sddlmZddlZddlZddlmZddlmZmZddl Z ddl m Z ddl m Z mZmZmZmZmZmZmZmZmZmZmZmZmZddlmZddlmZdd lm Z m!Z!m"Z"dd l#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*dd l+m,Z,m-Z-dd l.m/Z/d dl0m1Z1d dl2m3Z3erddl4m5Z5m6Z6gdZ7dde7DZ8gdZ9edddZ:dddZ;dd"d#Zdd(d)Z?dd*d+Z@dd,d-ZAdd.d/ZBdd0d1ZCdd2d3ZDd4d5dd=d>ZEddCdDZFd4d5ddKdLZGddMdNZHddOdPZId4d5ddSdTZJdUdUdVdd[d\ZKdd^d_ZLedd`daZMddcddZNddfdgZOddidjZPedkZQedlZRedmZSednZTd4d5ddodpZUddrdsZVedtZWd4d5ddydzZXdd|d}ZYdd~dZZd4d5dddZ[d4d5dddZ\d4d5dddZ]dS)) annotationsN) defaultdict)Callable TYPE_CHECKING)cpp)arg_parser_output_exprscpp_dispatch_exprscpp_dispatch_targetdispatch_lambda_argsdispatch_lambda_exprsdispatch_lambda_return_strhas_tensor_optionsPythonSignaturePythonSignatureDeprecatedPythonSignatureGroup!PythonSignatureNativeFunctionPair signaturesignature_from_schemastructseq_fieldnames) CodeTemplatewith_native_function) cpp_stringparse_native_yamlparse_tags_yaml)ArgumentBaseOperatorNameFunctionSchemaNativeFunction SchemaKindTypeVariant) FileManagersplit_name_params) YamlLoader)is_tensor_list_type) should_trace)IterableSequence)Salias contiguousZis_cudaZ is_sparseZ is_sparse_csrsizeZstrideZsym_sizeZ sym_strideZsym_storage_offsetZ sym_numelz .*_backwardz#.*_backward_(out|input|weight|bias)z .*_forwardz.*_forward_outz.*_jvpZ _unsafe_viewZtensorz2_?sparse_(coo|compressed|csr|csc|bsr|bsc)_tensor.*z_range.*Z_sparse_add_outz _sparse_div.*z _sparse_mul.*z _sparse_sub.*Z_sparse_dense_add_outindexZ index_outZunique_dim_consecutivez _cumsum.*z _cumprod.*z_sum.*z_prod.*z_th_.*z_thnn_.*zrange.*z_solve.*z _inverse.*z _cholesky.*z_triangular_solve.*z_qr.*z_svd.*sliceitemZ_local_scalar_densetoZ_to_copyZ _to_copy_outZ _reshape_copyZ_reshape_copy_outZcopy_sparse_to_sparse_Zcopy_Z _foreach_copyZnumpy_TZmatrix_HZmTZmHznonzero(_(out|numpy))?set_dataz.*_overrideabledataZis_leafZ output_nr_versionZrequires_grad_Z retains_gradset_Z _fw_primalZ)fake_quantize_per_tensor_affine_cachemaskZ*fake_quantize_per_channel_affine_cachemaskZ!_new_zeros_with_same_feature_metaZ_has_same_storage_numelZ_reshape_aliasZreplace_copyz fill.Tensorz fill.Scalarzlift.*Znormal_functionalnbytesitemsizeZ_batch_norm_with_updateZ_batch_norm_with_update_outZ_batch_norm_no_updatecCsg|] }td|dqS)^$)recompile).0patternr>^/var/www/auris/lib/python3.10/site-packages/torchgen/packaged/autograd/gen_python_functions.py sr@)z?add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> TensorzHadd_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)z?sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> TensorzHsub_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)z/mul.Scalar(Tensor self, Scalar other) -> Tensorz8mul_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)z/div.Scalar(Tensor self, Scalar other) -> Tensorz8div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)frreturnboolcCsdd|jvr d|jvr dSt|j}tD] }||rdSqt|j}tD] }||kr/dSq&dS)N generatedZ view_copyFT)tagsrnamefuncSKIP_PYTHON_BINDINGSmatchstrSKIP_PYTHON_BINDINGS_SIGNATURES)rArFZ skip_regexrr=r>r>r?should_generate_py_bindings   rLrFrrJcCs d|S)NZ THPVariable_r>rFr>r>r? get_pycname rN overloads+Sequence[PythonSignatureNativeFunctionPair]cCst|dko|djdkS)Nr%r)lenrarguments_count)rPr>r>r?is_noargsrTcC|jduo tj|jvSN) python_moduler!methodvariantsrAr>r>r?is_py_variable_methodr[cCrUrV)rWr!functionrYrZr>r>r?is_py_torch_functionr\r^cC |jdkS)NnnrWrZr>r>r?is_py_nn_functionrOrbcCr_)NZfftrarZr>r>r?is_py_fft_functionrOrccCr_)NZlinalgrarZr>r>r?is_py_linalg_functionrOrdcCr_)NnestedrarZr>r>r?is_py_nested_functionrOrfcCr_)NsparserarZr>r>r?is_py_sparse_functionrOrhcCr_)NZspecialrarZr>r>r?is_py_special_functionrOriTsymintoutnative_yaml_pathtags_yaml_pathdeprecated_yaml_path template_pathrkNonec s>t||dd}t||j}ttt|}t||dd}t||tddd|dt||dd} t || t dddd |d t|| t d d d|dt|| t d dd|dt|| t ddd|dt|| tddddt|| tddd|dt|| tddd|dt|| dddt|| dddt|d!fdd } |d | dS)"NF)Z install_dirZ template_dirdry_runTrXzpython_variable_methods.cpprXrktorchzpython_torch_functions.cpp)rX num_shardsrktorch.nnzpython_nn_functions.cpp torch.fftzpython_fft_functions.cpp torch.linalgzpython_linalg_functions.cpp torch.nestedzpython_nested_functions.cpp torch.sparsezpython_sparse_functions.cpp torch.specialzpython_special_functions.cppcSdSNTr>fnr>r>r?qzgen..zpython_return_types.cppcSr~rr>rr>r>r?rtrzpython_return_types.hrBdict[str, str]csddddtDiS)NZenum_of_valid_tagscSsg|] }d|d|dqS)z .value("z ", at::Tag::)r>)r<tagr>r>r?r@}sz.gen..gen_tags_enum..)joinsortedr> valid_tagsr>r? gen_tags_enumys zgen..gen_tags_enumzpython_enum_tag.cpp)rBr)r"rnative_functionslistfilterrLload_signaturescreate_python_bindingsr[create_python_bindings_shardedr^rbrcrdrfrhri"create_python_return_type_bindings)create_python_return_type_bindings_headerrwrite) rlrmrnrorprkfmrmethodsZ functionsrr>rr?gen s           rpairspred Callable[[NativeFunction], bool]?dict[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]cCs6tt}|D]}||jr||jjjj|q|SrV)rrr]rGrFappend)rrgroupedpairr>r>r?group_filter_overloadss rrr"module str | NonefilenamerXc sggggt||}t|tdD]1}||} t||| ||dt||| |dt|| |dd|j dq fdddS) +Generates Python bindings to ATen functionskeyrtrs#include cs$ddddS)N@generated from /)generated_comment ops_headers py_forwards py_methodspy_method_defstemplate_dir_for_commentsr>rrrrrrr>r?rsz(create_python_bindings..N) rrkeysrJr method_impl method_defextend forward_declsbasewrite_with_template) rrrrrrXrkrrFrPr>rr?rs$  rc sggt||}t|tdD]$}||}t|\}}|s#dnd||s/dnd|qfdddS)z Generate function to initialize and return named tuple for native functions which returns named tuple and registration invocations in `python_return_types.cpp`. rr cs ddddS)Nrrr)rZpy_return_typespy_return_types_registrationsrr>rrZpy_return_types_definitionrr>r?rs z4create_python_return_type_bindings..N)rrrrJ1generate_return_type_definition_and_registrationsrrr) rrrrrrFrP definitions registrationsr>rr?rs&  rcsjgt||}t|tdD]}||}t|}|sdnd|qfdddS)z Generate function to initialize and return named tuple for native functions which returns named tuple and relevant entry for the map in `python_return_types.cpp`. rrrcsddddS)Nrrr)rpy_return_types_declarationsrr>rrrr>r?rs z;create_python_return_type_bindings_header..N)rrrrJ!generate_return_type_declarationsrrr)rrrrrrFrP declarationsr>rr?rs rrwintc sbt||}ddd} dfdd } |j||d d d |d |i| | |hdddS)rkv@tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]rBrJcSs |djS)Nr)r)rr>r>r?key_funcs z0create_python_bindings_sharded..key_funcdict[str, list[str]]c sN|\}}d|jdgtt||dt||dgt||dgdS)Nrrrsrt)rrrr)rrrrr)rrFZfn_pairsrXrrkr>r?env_funcsz0create_python_bindings_sharded..env_funcrrrr>rrrr)Zbase_envZkey_fnZ env_callablerwZ sharded_keysN)rrrBrJ)rrrBr)rZ write_shardeditemsr) rrrrrrXrwrkrrrr>rr?rs   rF)skip_deprecatedpyirlist[NativeFunction]rrcsBtdfdd }tt||}t||d}|r|S||S) NrArrBrcstt|d|dS)NrXrrr])rrrZrr>r?gen_signature_pairs8s z,load_signatures..gen_signature_pairsr)rArrBr)rrmapload_deprecated_signatures)rrorXrrrr deprecatedr>rr?r0sr'list[PythonSignatureNativeFunctionPair]cstt}|D] }||jj|qg}t|}tj|td}Wdn1s*wY|D]} t | dt | d\} | drO| dd} dt didd jjDD]} | vsr| vsrJd | qadfdd } d} || D]3}| |jjsqd} t|jj||d}|tt|j|j|j|j|j|jt|jd |jdq| sJd| dtq1|S)N)LoaderrFZatenZ_outr1ScalarcSsi|]}|j|qSr>rMr<ar>r>r? iz.load_deprecated_signatures..z*deprecation definiton: Unrecognized value aten_schemarrBrCcsr t|jj|jj}n|jj}t|D]9\}}|tkrF|}|vr.|}d}n |}|j}|j }||jksB||j krEdSq|j durNdSqtj t|j koft ddt j |j DS)NFcss|] \}}||kVqdSrVr>)r<rbr>r>r? s  zKload_deprecated_signatures..is_schema_compatible..) itertoolschain argumentsrlZ flat_non_outflat_all enumeraterRtype annotationdefaultreturnsallzip)rriargZarg_nameZ schema_typeZschema_annotationZ schema_argZ call_argsZis_outZknown_constantsZschemaZschema_args_by_namer>r?is_schema_compatibleqs.    z8load_deprecated_signatures..is_schema_compatibleFT)category_overriderXr) rF input_args input_kwargs output_argstensor_options_argsrXdeprecated_schemaZdeprecated_args_exprsrrzNo native function with name z matched signature: )rrrBrC)rrrrFropenyamlloadr$rparser#endswithreplacer rrr]rGrrrrrrrrrXtuplerrJ)rrorXrrrresultsrAZdeprecated_defsrZ aten_namerFrZany_schema_foundZ python_sigr>rr?rFsj        rcCs(t|j}t|jj}d|g|S)N_)rrFrGrrr)rArF fieldnamesr>r>r?gen_structseq_typename_keys  r tuple[list[str], dict[str, str]]cCsi}g}|D]<}t|jjj}|sqt|jj}t|j}||}|durBd|s,dnt|}|||<| d|d|dq||fS)zr Generate block of named tuple type def inits, and add typeref snippets to declarations that use them N NamedTuplerzstatic PyTypeObject* z = generated::get_ _structseq(); rr]rGrrrFrgetrRr)rP typenamesZtypedefsoverloadrrFtn_keytypenamer>r>r?emit_structseq_calls,   r tuple[list[str], list[str]]c Csi}g}g}|D]h}t|jjj}|sqddd|D}t|jj}t|j}||} | durp|d|s:dnt |} | ||<| d|d|d | d |d t |d | d | d| d| d|d|dq||fS)z Generate block of function in `python_return_types.cpp` to initialize and return named tuple for a native function which returns named tuple and registration invocations in same file. , css|] }d|dVqdS){"z", ""}Nr>)r<rr>r>r?rszDgenerate_return_type_definition_and_registrations..NrrPyTypeObject* get_zI_structseq() { static PyStructSequence_Field NamedTuple_fields[] = { z(, {nullptr} }; static PyTypeObject zh; static bool is_initialized = false; static PyStructSequence_Desc desc = { "torch.return_types.z", nullptr, NamedTuple_fields, zB }; if (!is_initialized) { PyStructSequence_InitType(&z, &desc); zm.tp_repr = (reprfunc)torch::utils::returned_structseq_repr; is_initialized = true; } return &z; } z$addReturnType(return_types_module, "z", generated::get_z_structseq());) rr]rGrrrrFrrrRr) rPrrrrrfieldsrFrr r>r>r?rsN    r list[str]cCsi}g}|D];}t|jjj}|sqt|jj}t|j}||}|durA|d|s.dnt|}|||<| d|dq|S)z Generate block of function declarations in `python_return_types.h` to initialize and return named tuple for a native function. Nrrrrr)rPrrrrrFrr r>r>r?rs"   ra\ // ${name} static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { ${method_header} static PythonArgParser parser({ ${signatures} }, /*traceable=*/${traceable}); ParsedArgs<${max_args}> parsed_args; auto _r = parser.parse(${self_}, args, kwargs, parsed_args); ${check_has_torch_function} switch (_r.idx) { ${dispatch} } ${method_footer} } z&case ${overload_index}: { ${body} } ao// ${name} static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { ${method_header} static PythonArgParser parser({ ${signatures} }, /*traceable=*/${traceable}); ParsedArgs<${max_args}> parsed_args; auto _r = parser.parse(${self_}, args, kwargs, parsed_args); ${check_has_torch_function} ${dispatch} ${method_footer} } z// ${name} static PyObject * ${pycname}(PyObject* self_, PyObject* args) { ${method_header} ${check_has_torch_function} ${dispatch} ${method_footer} } c CsHt|}t|}t|\}}dg} | |7} | |rdgng7} |r"gndgdg} tdd|Dr3dnd} t||d } t| d k} g}g}t| D],\}}|jj|d }| t t |d t |||d }| | srt j||d n|qI|r{t}n| rt}nt}|j||| td d|D|| t||||d|| |rdd Sdd S)z? Generate a python binding for all overloads of an op. ZHANDLE_TH_ERRORSz/const Tensor& self = THPVariable_Unpack(self_);Py_RETURN_NONE;ZEND_HANDLE_TH_ERRORScss|]}t|jVqdSrV)r'r]r<or>r>r?rzmethod_impl..truefalserjr%,)overload_indexbodycss|]}|jVqdSrV)rrSrr>r>r?rr)rFrnoargrXself_nullptr) rFpycname method_headerZmax_args signatures traceableZcheck_has_torch_functiondispatch method_footerr)rNrTr rgroup_overloadsrRrr signature_strrrrJemit_dispatch_casePY_VARIABLE_CASE substitutePY_VARIABLE_METHOD_NOARGS$PY_VARIABLE_METHOD_VARARGS_SINGLETONPY_VARIABLE_METHOD_VARARGSmaxgen_has_torch_function_check)rFrrPrXrkrrZstructseq_initsstructseq_typenamesrr"r grouped_overloadsZ is_singletonrr!rrrZ dispatch_bodytemplater>r>r?rsf    rrcCs`|r |r d|dSdS|rdnd}|r dddd d d d d |nd}d|d|d|p,ddS)NzMif(check_has_torch_function(self_)) { return handle_torch_function(self_, "z"); } rrrZTHPVariableFunctionsModuleZTHPNNVariableFunctionsModuleZTHPFFTVariableFunctionsModuleZ THPLinalgVariableFunctionsModuleZ THPNestedVariableFunctionsModuleZ THPSparseVariableFunctionsModuleZ!THPSpecialVariableFunctionsModule)rurxryrzr{r|r}ZTHPVariableClasszAif(_r.has_torch_function()) { return handle_torch_function(_r, z, args, kwargs, z, "z torch.Tensorr>)rFrrrXr namespacer>r>r?r,s8   r,zRif (_r.isNone(${out_idx})) { ${call_dispatch} } else { ${call_dispatch_out} } rrr-rc CsT|jdur tj|jt|j|j||dt|j|j||ddSt|j|j||dS)a0 Emit dispatch code for a single parsed signature. This corresponds to either a single native function, or a pair that differ only in output params. In the latter case, a single python signature is used for both and dispatching switches on the presence/absence of passed output args. Nrj)Zout_idxZ call_dispatchZcall_dispatch_out)ZoutplacePY_VARIABLE_OUTr'rZ output_idxemit_single_dispatchr)rr-rkr>r>r?r%s   r%tuple[str, ...]cCs4|rdSt|}t|rd|dfSd|dfS)Nr>zstatic PyObject * z#(PyObject* self_, PyObject* args); z5(PyObject* self_, PyObject* args, PyObject* kwargs); )rNrT)rFrPrXrr>r>r?rsrcCsht|}|jr d|d}t|r|rdnd}nd|d}d}|dkr(|d7}d |d |d |d S) z$ Generate method def entry. zTypeError_to_NotImplemented_<>Z METH_NOARGSzMETH_VARARGS | METH_KEYWORDSzcastPyCFunctionWithKeywords(rruz | METH_STATICr z", r z , nullptr},)rNZ dunder_methodrT)rFrrPrXrflagsr>r>r?r<s   rSequence[PythonSignatureGroup]c s`i}i|D]C}|jjd|d}|jjr0|vr+td|jjd|jjd||<q||vrEtd|jjd||jjd|||<qD]P\}}||vrg}|D]'}t|jjjjt|jjjjkr|jjs|jj s| |jjd|dqZ|jj|d}td|d|d d d d |DqNfd d|D}t ||dS)NT) skip_outputsrkz(Found duplicated function definition: - z. Existing definition: - .rjz4While identifying overloads, we found an out schema z] without a corresponding non-out variant. We expected the non-out variant to have schema: - zy Please check that you spelled the schema correctly in native_functions.yaml. We discovered the following candidate(s): rcss|]}d|VqdS)z- Nr>)r< candidater>r>r?rrz"group_overloads..cs$g|]\}}tj||dqS))Z functionalrl)rZ from_pairsr)r<sigrZ outplacesr>r?r@sz#group_overloads..) rr$r]rGZ is_out_fn RuntimeErrorrrJrFrrrsort_overloads) rPrkbasesrr:rl candidatesZout_sigrr>r;r?r#_sb          r#r.c sddddfd d }tfd ddtttD]\}}tD]\}}||j|jr8||q&qs@tSt}ttfddt |}t |D]"} || } t D]} | } | | | sv| =| | qaqUfdd|DS)Nt1r t2rBrCcSst|dkr t|dkpbt|dkot|dkpbdt|vo#dt|vpbt|dko5t|dkp5t|dkpbt|d koDt|d d kpbt|d koPt|dkpbt|d ks]t|dkobt|dkS)NrZTensorzScalar?zTensor?ZDimnamezint[]rzint?zTensor[]z[]zSymInt[]ZSymInt)rJfind)r@rAr>r>r?is_arg_smallers*           z&sort_overloads..is_arg_smallers1rs2csl|jdd|jdd}}t|t|krdStddt||D}tfddt||D}|o5| S)z-Returns True if s1 < s2 in the partial order.T)r7Fcss |] \}}|j|jkVqdSrV)rr<Zarg1Zarg2r>r>r?rsz5sort_overloads..is_smaller..c3s6|]\}}t|jt|jkp|j|jVqdSrV)rJrrGrDr>r?rs  )rrRrr)rErFZargs1Zargs2equalZsmaller_or_equalrHr>r? is_smallers  z"sort_overloads..is_smallercs|jjdS)Nrj)rr$xrjr>r?rsz sort_overloads..rcs|vSrVr>rK) larger_thanr>r?rscsg|]}|qSr>r>)r<rL)r.r>r?r@rz"sort_overloads..)r@r rAr rBrC)rErrFrrBrC) rrsetrraddrrRrrangerdiscardr) r.rkrJi1Z overload1i2Z overload2NZ sorted_idsidxrjZlargerr>)r.rDrMrkr?r=s6     r=psrcstdfdd }||S) z: Emit dispatch code for a single native function. rArrBrJcsttr dj}nd|j}t|j}dddt|dD}t|}t |}dt |d}t |d}t |d}d|j } d|j} jo`t| p`jo`d |v} | rld |d jd nd } |d kr|jjj} t|jjdr|jtjkr| durt| jjsJd}nd}|d| d|d|d|d|d|d|d| d | d|dSt|}|dur|dnd }|d| d|d|d|d|d|d|d|d| d | dS)Nz// [deprecated] aten::z // aten::r css"|] }|jd|jVqdS) N)Ztype_strrFrr>r>r?r2s z3emit_single_dispatch..go..rj)Zpython_signaturerZ requires_gradz.set_requires_grad(rrvoidZ _foreach_z\PyObject* self_tensorlist = _r.args[0]; Py_INCREF(self_tensorlist); return self_tensorlist; rz auto dispatch_z = [](z) -> z, { pybind11::gil_scoped_release no_gil; (z); }; dispatch_z; z3 { pybind11::gil_scoped_release no_gil; return z); }; return wrap(Z dispatch_z); ) isinstancerrrGrrFrr r r rrr initsexprsrr rXexprrself_argrJ startswithkindrZinplacer&argumentrrr)rAZschema_commentrFZlambda_formalsZ lambda_returnZdispatch_calleeZ dispatch_argsZparser_outputsZlambda_arg_exprsr\Z lambda_argsZneed_set_requires_gradZset_requires_gradr_Z return_stmtr Zstructseq_typerefrWr-rkr>r?go(s            z emit_single_dispatch..goNrArrBrJr)rWrAr-rkrdr>rcr?r2s Sr2)rArrBrC)rFrrBrJ)rPrQrBrC)rlrJrmrJrnrJrorJrprJrkrCrBrq)rrQrrrBr)rr"rrQrrrrrrJrXrCrkrCrBrq) rr"rrQrrrrJrBrq)rr"rrQrrrrrrJrXrCrwrrkrCrBrq) rrrorJrXrCrrCrrCrBrQ) rrQrorJrXrCrrCrBrre)rPrQrBr)rPrQrBr )rPrQrBr) rFrrrrPrQrXrCrkrCrBrJ) rFrrrrrCrXrCrBrJ)rrr-rrkrCrBrJ)rFrrPrQrXrCrBr3) rFrrrrPrQrXrCrBrJ)rPrQrkrCrBr6)r.r6rkrCrBr6) rWrrArr-rrkrCrBrJ)^ __future__rrr: collectionsrtypingrrrZ torchgen.apirZtorchgen.api.pythonrrr r r r r rrrrrrrZtorchgen.code_templaterZtorchgen.contextrZ torchgen.genrrrZtorchgen.modelrrrrrr r!Ztorchgen.utilsr"r#Ztorchgen.yaml_utilsr$Zgen_inplace_or_view_typer&Zgen_trace_typer'collections.abcr(r)Z_SKIP_PYTHON_BINDINGSrHrKrLrNrTr[r^rbrcrdrfrhrirrrrrrrrrr rrr*r&r)r(rr,r1r%rrr#r=r2r>r>r>r?s !  @  $    V            | ) ')3 v    3&  E$ ( $g]