o he"@sddlZddlmZddlmZmZmZmZddlmZm Z m Z m Z ddl m Z ddlmZGdd d eZejGd d d Zd e jd edefddZdedefddZddZddZe jZe jZe jZdS)N) dataclass)AnyCallableOptionalProtocol)_C_opsautogradTensor)_pytree)utilsc@&eZdZUeeed<eeed<dS) InfoProtocol _backward_fn_setup_context_fnN__name__ __module__ __qualname__rr__annotations__rrk/var/www/html/construction_image-detection-poc/venv/lib/python3.10/site-packages/torch/_library/autograd.pyr s rc@r)InforrNrrrrrrs ropinforeturncsdjdjdj}tjtGdddfddfdd}fd d }t|tj ft |t |d j}t d d g|j |j RDrXtfdd}|S)NGeneratedBackwardFor__c@s(eZdZUejed<eeefed<dS)z$make_autograd_impl..Metadatakeysetkeyword_only_argsN) rrrrDispatchKeySetrdictstrrrrrrMetadatas r$csn|d}|dd}t|j}|j}j|tj@g|Ri|}|WdS1s0wYdS)N)r_AutoDispatchBelowAutogradrr redispatch_after_autograd_keyset)argsmetadatarkwargsresult)rrrforward_no_grad!s  $z+make_autograd_impl..forward_no_gradcs|d}|dd}t@|j}|j}j|tj@g|Ri|}jrEtj ||\}}r=j||||dnj|||d|WdS1sQwYdS)Nr%)ctxinputskeyword_only_inputsoutput)r.r/r1) rr&rr r'r(rr fill_defaults_schema)r.r)r*rr+r,)has_kwarg_only_argsrrrrforward+s   $z#make_autograd_impl..forwardcsrjr1z|j}|jdd|_j|g|R}W||_n||_wt|tr-g|dRS|dfStdd)Nr%zTrying to backward through zQ but no autograd formula was registered. Please use register_autograd to add one.)rneeds_input_grad isinstancetuple RuntimeError)r.gradsprev_needs_input_gradr,)rrrrbackwardHs  z$make_autograd_impl..backward)r5r<css|] }t|jVqdSN)r is_tensorlist_like_typetype).0arrr cs  z%make_autograd_impl..csVtrttdd|trjg|||R}|Sg|||R}|S)NcSs|jSr=) requires_grad)xrrrmsz;make_autograd_impl..autograd_impl..)ris_grad_enabledr tree_any_onlyr not_list_of_tensorapply)rr)r r,) Generatedr$r-rr autograd_implks  z)make_autograd_impl..autograd_impl) _namespace_opname _overloadnamer r4r3rr?r Function staticmethodany argumentsreturnssupports_tensorlist)rrnamer5r<schemarKr)rJr$r-r4rrrmake_autograd_impls*     rWclscsb|j|j|jtGdddfdd}fdd}fdd}||_||_||_|S) a4Allows a given autograd.Function class to support List[Tensor] inputs/outputs. Regular autograd.Function has a constraint that it only directly supports autograd for Tensors. Applying @supports_tensorlist enables an autograd.Function to support autograd for List[Tensor] inputs and outputs. c@s6eZdZUeed<dZeeed<dZeeed<dS)z%supports_tensorlist..Metadata input_specN output_specresult_is_tuple) rrrspec_trrZrr[boolrrrrr$s r$cs|d}|dd}t|stdtt||j}|g|R}t|t|_|js/|f}t|t\}}||_ t |drBt d||_ t|S)Nr%zNYI: calling supports_tensorlist autograd.Function.forward directly. You should probably be calling .apply instead. Please file an issue if not. _pt_metadataz@Please don't set ctx._pt_metadata; PyTorch uses it to store info) r7NotImplementedError unflattenlistrYr8r[flattenrHrZhasattrr9r^)r.r)r*r, flat_resultrZ)r$ orig_forwardrr new_forwards&    z(supports_tensorlist..new_forwardcst|ds td|j}tt||j}|j}ztt|jdd|j|_|g|R}W||_n||_wt|t s?|f}t |t \}}||jkrWt d|d|jdt |dgS)Nr^zNYI: calling supports_tensorlist autograd.Function.backward directly. This will automatically get called by PyTorch autograd. Please file an issue if you need this.r%zRExpected the return from backward to be of the same structure as the inputs. Got: z (return from backward), z (inputs)) rcr_r^r`rarZr6rYr7r8rbnot_list_of_optional_tensorr9)r.r:r*r; grad_inputsflat_grad_inputsgrad_inputs_spec) orig_backwardrr new_backwards6   z)supports_tensorlist..new_backwardcsxt|td\}}|}g||R}|jdusJtt||j}|js:t|ts.Jt|dks6J|dS|S)N)is_leafr r) rbrHrZr`rar[r7r8len)r) flat_argsrYr*r,)r$ orig_applyrr new_applysz&supports_tensorlist..new_apply)r5r<rIr)rXrfrlrqr)r$rprkrerrTws , rTcC.t|trdSt|trtdd|DSdS)NFcss|] }t|t VqdSr=r7r r@lrrrrBsz%not_list_of_tensor..Tr7r8rarQtreerrrrH  rHcCrr)NFcss$|] }|duo t|t VqdSr=rsrtrrrrBs"z.not_list_of_optional_tensor..Trvrwrrrrgryrg) dataclassesrtypingrrrrtorchrrr r torch.utilsr r rr OpOverloadrWrTrHrg tree_flattenrbtree_unflattenr`TreeSpecr\rrrrs    `h