+ j( ^RIt^RIHt^RIHt^RIHt^RIHtH t H t ^RI H t H t ^RIt^RIHt^RIHt^RIHt] !R R R 7t] !R 4tRR lt]!4t]!4t!RR4t] !RRR R7t] !R]P8]]P8,4t!RR] ] ]]]3,4t]3RRllt RRlt!R#)N) OrderedDict)Callablewraps) ConcatenateGenericProtocol) ParamSpecTypeVar)_State)_get_root_modules_TT) covariant_PcHV R\\P!44 2#)_)struuiduuid4)strings&g/home/luca/screen-checking/.venv/lib/python3.14/site-packages/torch/distributed/_composable/contract.pygenerate_state_keyrsXQs4::<() **c]tRt^tRtR#) RegistryItemN)__name__ __module__ __qualname____firstlineno____static_attributes__rrrrrsrr_TStater )boundr_McDa]tRt^'toV3RlRltV3RlRltRtVtR#) _ContractFncT<V^8dQhRS[PRS[PRS[/#)argskwargsreturn)rr)r*r)format __classdict__s"r __annotate___ContractFn.__annotate__(s FFbggFFrFrcR#Nr)selfr)r*s&*,r__call___ContractFn.__call__(s3rc:<V^8dQhRS[PRS[/#r(moduler+)nnModuler")r,r-s"rr.r/*s66BII6'6rcR#r1r)r2r7s&&rstate_ContractFn.state*s3rrN)rrrr r3r;r!__classdictcell__)r-s@rr&r&'sFF66rr&c V^8dQhR\\,R\\\\\ 3,\3,.\ \\\ 3,\\3,3,/#)r( state_clsr+)typer"rrr$rr&)r,s"rr.r.-s_DDG}D  k"b&!2%&' BF#R013Drc6a\S4RV3Rll4pV#)a Decorate a function as a composable distributed API, where the first argument of the function must be an :class:`nn.Module` instance or sequence of :class:`nn.Module` instances. The decorator verifies that the decorated function does not modify fully-qualified names (FQNs) for parameters, buffers, or modules. The decorated function can return different module instances than the input modules; the FQN invariant will be enforced following the input order. When a function ``func`` is decorated by ``@contract()``, a ``.state(module: nn.Module)`` method will be installed to the decorated function. Then you can retrieve and modify the state on a module by calling ``func.state(module)``. Example:: >>> # xdoctest: +SKIP >>> import torch.nn as nn >>> >>> class MyModel(nn.Module): >>> def __init__(self) -> None: >>> super().__init__() >>> self.l1 = nn.Linear(10, 10) >>> self.l2 = nn.Linear(10, 10) >>> >>> def forward(self, x): >>> return self.l2(self.l1(x)) >>> >>> @contract() >>> def my_feature(module: nn.Module) -> nn.Module: >>> my_feature.state(module).some_state = "any value" >>> return module >>> >>> model = MyModel() >>> my_feature(model.l1) >>> assert my_feature.state(model.l1).some_state == "any value" >>> my_feature(model.l2) >>> model(torch.randn(2, 10)).sum().backward() cV^8dQhR\\\\3,\3,R\\\\3,\\ 3,/#)r(funcr+)rrr$rr&r")r,s"rr.contract..__annotate__]sHRR{2r6*B./R [R("g5 6RrcX<a\S4RVV3Rll4pRV3RllpW!nV#)cpV^8dQhR\R\PR\PR\/#)r(r7r)r*r+)r$rr)r*)r,s"rr.-contract..inner..__annotate__as>D D D 77D iiD  D rc<Tp\V\P4'dV.pM\\ V44pS!4p\ 4p.p.p.p VEFyp\ 4p \ 4p VPP\V 4p \V \4'g\RV 24hVPP\V 4p \V \4'g\RV 24hSV 9gSPV 9d\RSP RV 24hV PSV4V PSPV4VP\ VP444VP\ VP!444V P\ VP#444EK| S!V.VO5/VBpVfTp\V\P4'dV.pM\\ V44p.p.p.pVF{pVP\ VP444VP\ VP!444VP\ VP#444K} \%V 4p\%V4pVV8wd\SP RV RV R24hRRlp\'VV4F>wppV!\ VP)44\ VP)44R 4K@ \'VV4F>wppV!\ VP)44\ VP)44R 4K@ \'V V4F>wppV!\ VP)44\ VP)44R 4K@ V#) z-Distributed composable API states corrupted: z/Distributed composable API registry corrupted: zOEach distinct composable distributed API can only be applied to a module once. z3 has already been applied to the following module: zB should return the same number of modules as input modulesInputs: z modules Outputs: z moduleschV^8dQhR\\,R\\,R\/#)r( orig_fqnsnew_fqns check_key)listr)r,s"rr.>contract..inner..wrapper..__annotate__s*  T#Y $s) PS rcW8XdR#\V4\V4rCW4, pWC, p\V4'g\V4'd\V RV RV 24h\V RV RV 24h)NzVComposable distributed API implementations cannot modify FQNs. FQNs only in original: z FQNs only in new: z[Composable distributed API implementations cannot modify the order of FQNs. Original FQNs: z New FQNs: )setlen RuntimeError)rJrKrL orig_fqn_set new_fqn_set orig_onlynew_onlys&&& r check_fqn;contract..inner..wrapper..check_fqns(,/ NCMk(6 &5y>>S]]&$+22;=--5J8'$+**35%%-J 0rzChecking parameters: zChecking buffers: zChecking modules: ) isinstancer8r9r rMrr__dict__ setdefault STATE_KEYdictAssertionError REGISTRY_KEYrappendnamed_parameters named_buffers named_modulesrQzipkeys)r7r)r* inp_modulemodulesr; registry_itemall_orig_named_paramsall_orig_named_buffersall_orig_named_modulesdefault_all_statedefault_registry all_stateregistryupdatedupdated_modulesall_new_named_paramsall_new_named_buffersall_new_named_modulesnum_orig_modulesnum_new_modulesrWorig_named_paramsnew_named_paramsorig_named_buffersnew_named_buffersorig_named_modulesnew_named_modulesrCr?s&*, rwrapper(contract..inner..wrapper`s  J&")),,!( ,DL9KE(NM DF !DF "AC "!.inner..get_states+??--c$i r)rr;)rCr}rr?sf rinnercontract..inner\s: tD D  D L   " rr)r?rsf rcontractr-s'^ 9RRh LrctV^8dQhR\PR\\\3,R,/#)r(r7r+N)r8r9r]rr)r,s"rr.r.s-//"))/S,->(?$(F/rc$\V\R4#)z Get an ``OrderedDict`` of composable APIs that have been applied to the ``module``, indexed by the API name. If no API has been applied, then this returns ``None``. N)getattrr_)r7s&r _get_registryrs 6< ..r)__composable_api_state_key)"r collectionsrcollections.abcr functoolsrtypingrrr typing_extensionsr r torchtorch.nnr8#torch.distributed._composable_stater torch.distributed.utilsr rrrr\r_rr"r9rMr$r&rrrrrrs #$110 65 TT"t_+  !#    )8t < T299d299o.7(GBGO47 &DN/r