o 8} i@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_viewtensorz2_?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.*sliceitem_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_nrZ_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_meta_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@E/home/app/PyTorch/test/pytorch/tools/autograd/gen_python_functions.py srB)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 generated view_copyFT)tagsrnamefuncSKIP_PYTHON_BINDINGSmatchstrSKIP_PYTHON_BINDINGS_SIGNATURES)rCrIZ skip_regexrr?r@r@rAshould_generate_py_bindings   rOrIrrMcCs d|S)NZ THPVariable_r@rIr@r@rA get_pycname rQ overloads+Sequence[PythonSignatureNativeFunctionPair]cCst|dko|djdkS)Nr%r)lenrarguments_count)rSr@r@rAis_noargsrWcC|jduo tj|jvSN) python_moduler!methodvariantsrCr@r@rAis_py_variable_methodr^cCrXrY)rZr!functionr\r]r@r@rAis_py_torch_functionr_racC |jdkS)NnnrZr]r@r@rAis_py_nn_functionrRrecCrb)NZfftrdr]r@r@rAis_py_fft_functionrRrfcCrb)NZlinalgrdr]r@r@rAis_py_linalg_functionrRrgcCrb)Nnestedrdr]r@r@rAis_py_nested_functionrRricCrb)NZsparserdr]r@r@rAis_py_sparse_functionrRrjcCrb)NZspecialrdr]r@r@rAis_py_special_functionrRrkTsymintoutnative_yaml_pathtags_yaml_pathdeprecated_yaml_path template_pathrmNonec 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) install_dir template_dirdry_runTr[zpython_variable_methods.cppr[rmtorchzpython_torch_functions.cpp)r[ num_shardsrmtorch.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@rAqzgen..zpython_return_types.cppcSrrr@rr@r@rArtrzpython_return_types.hrDdict[str, str]csddddtDiS)Nenum_of_valid_tagscSsg|] }d|d|dqS)z .value("z ", at::Tag::)r@)r>tagr@r@rArB}sz.gen..gen_tags_enum..)joinsortedr@ valid_tagsr@rA gen_tags_enumys zgen..gen_tags_enumzpython_enum_tag.cpp)rDr)r"rnative_functionslistfilterrOload_signaturescreate_python_bindingsr^create_python_bindings_shardedrarerfrgrirjrk"create_python_return_type_bindings)create_python_return_type_bindings_headerrwrite) rnrorprqrrrmfmrmethods functionsrr@rrAgen s           rpairspred Callable[[NativeFunction], bool]?dict[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]cCs6tt}|D]}||jr||jjjj|q|SrY)rrr`rJrIappend)rrgroupedpairr@r@rAgroup_filter_overloadss rrr"module str | Nonefilenamer[c sggggt||}t|tdD]1}||} t||| ||dt||| |dt|| |dd|j dq fdddS) +Generates Python bindings to ATen functionskeyrxrw#include cs$ddddS)N@generated from /)generated_comment ops_headers py_forwards py_methodspy_method_defstemplate_dir_for_commentsr@rrrrrrr@rArsz(create_python_bindings..N) rrkeysrMr method_impl method_defextend forward_declsbasewrite_with_template) rrrrrr[rmrrIrSr@rrArs$  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@rArs z4create_python_return_type_bindings..N)rrrrM1generate_return_type_definition_and_registrationsrrr) rrrrrrIrS definitions registrationsr@rrArs&  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@rArs z;create_python_return_type_bindings_header..N)rrrrM!generate_return_type_declarationsrrr)rrrrrrIrS declarationsr@rrArs rr{intc sbt||}ddd} dfdd } |j||d d d |d |i| | |hdddS)rkv@tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]rDrMcSs |djS)Nr)r)rr@r@rAkey_funcs z0create_python_bindings_sharded..key_funcdict[str, list[str]]c sN|\}}d|jdgtt||dt||dgt||dgdS)Nrrrwrx)rrrr)rrrrr)rrIZfn_pairsr[rrmr@rAenv_funcsz0create_python_bindings_sharded..env_funcrrrr>rrrr)base_envkey_fn env_callabler{ sharded_keysN)rrrDrM)rrrDr)r write_shardeditemsr) rrrrrr[r{rmrrrr@rrArs   rF)skip_deprecatedpyirlist[NativeFunction]rrcsBtdfdd }tt||}t||d}|r|S||S) NrCrrDrcstt|d|dS)Nr[rrr`)rrr]rr@rAgen_signature_pairs8s z,load_signatures..gen_signature_pairsr)rCrrDr)rrmapload_deprecated_signatures)rrqr[rrrr deprecatedr@rrAr0sr'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)LoaderrIaten_outr1ScalarcSsi|]}|j|qSr@rPr>ar@r@rA iz.load_deprecated_signatures..z*deprecation definiton: Unrecognized value aten_schemarrDrEcsr 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|] \}}||kVqdSrYr@)r>rbr@r@rA s  zKload_deprecated_signatures..is_schema_compatible..) itertoolschain argumentsrn flat_non_outflat_all enumeraterUtype annotationdefaultreturnsallzip)rriargZarg_nameZ schema_typeZschema_annotationZ schema_argZ call_argsis_outZknown_constantsschemaZschema_args_by_namer@rAis_schema_compatibleqs.    z8load_deprecated_signatures..is_schema_compatibleFT)category_overrider[r) rI input_args input_kwargs output_argstensor_options_argsr[deprecated_schemadeprecated_args_exprsrrzNo native function with name z matched signature: )rrrDrE)rrrrIropenyamlloadr$rparser#endswithreplacer rrr`rJrrrrrrrrr[tuplerrM)rrqr[rrrresultsrCZdeprecated_defsr aten_namerIrZany_schema_foundZ python_sigr@rrArFsj        rcCs(t|j}t|jj}d|g|S)N_)rrIrJrrr)rCrI fieldnamesr@r@rAgen_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`rJrrrIrgetrUr)rS typenamesZtypedefsoverloadrrItn_keytypenamer@r@rAemit_structseq_calls,   rtuple[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@rArszDgenerate_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`rJrrrrIrrrUr) rSrrrrrfieldsrIrrr@r@rArsN    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. Nrrr rr)rSrrrrrIrrr@r@rArs"   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|jVqdSrY)r'r`r>or@r@rArzmethod_impl..truefalserlr%,)overload_indexbodycss|]}|jVqdSrY)rrVr$r@r@rArr&)rIrnoargr[self_nullptr) rIpycname method_headerZmax_args signatures traceableZcheck_has_torch_functiondispatch method_footerr-)rQrWrrgroup_overloadsrUrr signature_strrrrMemit_dispatch_casePY_VARIABLE_CASE substitutePY_VARIABLE_METHOD_NOARGS$PY_VARIABLE_METHOD_VARARGS_SINGLETONPY_VARIABLE_METHOD_VARARGSmaxgen_has_torch_function_check)rIrrSr[rmr/r,Zstructseq_initsstructseq_typenamesr0r4r2grouped_overloadsZ is_singletonr1r3r*rrZ dispatch_bodytemplater@r@rArsf    rr,cCs`|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"); } rr-r.ZTHPVariableFunctionsModuleZTHPNNVariableFunctionsModuleZTHPFFTVariableFunctionsModuleZ THPLinalgVariableFunctionsModuleZ THPNestedVariableFunctionsModuleZ THPSparseVariableFunctionsModuleZ!THPSpecialVariableFunctionsModule)ryr|r}r~rrrZTHPVariableClasszAif(_r.has_torch_function()) { return handle_torch_function(_r, z, args, kwargs, z, "z torch.Tensorr@)rIrr,r[r- namespacer@r@rAr>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. Nrl)Zout_idxZ call_dispatchZcall_dispatch_out)outplacePY_VARIABLE_OUTr9r output_idxemit_single_dispatchr)rr?rmr@r@rAr7s   r7tuple[str, ...]cCs4|rdSt|}t|rd|dfSd|dfS)Nr@zstatic PyObject * z#(PyObject* self_, PyObject* args); z5(PyObject* self_, PyObject* args, PyObject* kwargs); )rQrW)rIrSr[r/r@r@rArsrcCsht|}|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(rryz | METH_STATICrz", rz , nullptr},)rQ dunder_methodrW)rIrrSr[r/flagsr@r@rAr<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_outputsrmz(Found duplicated function definition: - z. Existing definition: - .rlz4While 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@rArr&z"group_overloads..cs$g|]\}}tj||dqS)) functionalrn)r from_pairsr)r>sigrZ outplacesr@rArBsz#group_overloads..) rr6r`rJ is_out_fn RuntimeErrorrrMrIrrrsort_overloads) rSrmbasesrrQrnZ candidatesout_sigrr@rRrAr5_sb          r5r@c sddddfd d }tfd ddtttD]\}}tD]\}}||j|jr8||q&qs@tSt}ttfddt |}t |D]"} || } t D]} | } | | | sv| =| | qaqUfdd|DS)Nt1r t2rDrEcSst|dkr t|dkpbt|dkot|dkpbdt|vo#dt|vpbt|dko5t|dkp5t|dkpbt|d koDt|d d kpbt|d koPt|dkpbt|d ks]t|dkobt|dkS)NrTensorzScalar?zTensor?Dimnamezint[]rzint?zTensor[]z[]zSymInt[]SymInt)rMfind)rXrYr@r@rAis_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)rLFcss |] \}}|j|jkVqdSrY)rr>Zarg1Zarg2r@r@rArsz5sort_overloads..is_smaller..c3s6|]\}}t|jt|jkp|j|jVqdSrY)rMrrbr_r@rArs  )rrUrr)r`raZargs1Zargs2equalZsmaller_or_equalrcr@rA is_smallers  z"sort_overloads..is_smallercs|jjdS)Nrl)rr6xrlr@rArsz sort_overloads..rcs|vSrYr@rf) larger_thanr@rArscsg|]}|qSr@r@)r>rg)r@r@rArBrz"sort_overloads..)rXr rYr rDrE)r`rrarrDrE) rrsetrraddrrUrrangerdiscardr) r@rmreZi1Z overload1Zi2Z overload2NZ sorted_idsidxrjZlargerr@)r@r_rhrmrArUs6     rUpsrcstdfdd }||S) z: Emit dispatch code for a single native function. rCrrDrMcsttr 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::rcss"|] }|jd|jVqdS) N)type_strrIrr@r@rAr2s z3emit_single_dispatch..go..rl)python_signaturer requires_gradz.set_requires_grad(rrvoidZ _foreach_z\PyObject* self_tensorlist = _r.args[0]; Py_INCREF(self_tensorlist); return self_tensorlist; r#z 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); ) isinstancerrrJrrIrr r r rrr initsexprsrr r[exprrself_argrM startswithkindrinplacer&argumentrrr)rCZschema_commentrIZlambda_formalsZ lambda_returnZdispatch_calleeZ dispatch_argsZparser_outputsZlambda_arg_exprsrx lambda_argsZneed_set_requires_gradZset_requires_gradr{Z return_stmtrZstructseq_typerefrpr?rmr@rAgo(s            z emit_single_dispatch..goNrCrrDrMr)rprCr?rmrr@rrArFs SrF)rCrrDrE)rIrrDrM)rSrTrDrE)rnrMrorMrprMrqrMrrrMrmrErDrs)rrTrrrDr)rr"rrTrrrrrrMr[rErmrErDrs) rr"rrTrrrrMrDrs)rr"rrTrrrrrrMr[rEr{rrmrErDrs) rrrqrMr[rErrErrErDrT) rrTrqrMr[rErrErDrr)rSrTrDr)rSrTrDr)rSrTrDr") rIrrrrSrTr[rErmrErDrM) rIrrrr,rEr[rErDrM)rrr?rrmrErDrM)rIrrSrTr[rErDrG) rIrrrrSrTr[rErDrM)rSrTrmrErDrK)r@rKrmrErDrK) rprrCrr?rrmrErDrM)^ __future__rrr< collectionsrtypingrrr torchgen.apirtorchgen.api.pythonrrr r r r r rrrrrrrtorchgen.code_templatertorchgen.contextr torchgen.genrrrtorchgen.modelrrrrrr r!torchgen.utilsr"r#torchgen.yaml_utilsr$Zgen_inplace_or_view_typer&Zgen_trace_typer'collections.abcr(r)Z_SKIP_PYTHON_BINDINGSrKrNrOrQrWr^rarerfrgrirjrkrrrrrrrrrrrrr<r8r;r:rr>rDr7rrr5rUrFr@r@r@rAs !  @  $    V            | ) ')3 v    3&  E$ ( $g]