[Th( SSKrSSKJr SSKJr SSKJrJrJrJ r SSK J r J r J r SSKrSSKJr SSKJr SSKJr \ "SS S 9r\ "S 5rSS jr\"5r\"5r"S S5r\ "SSS S9r\ "S\R8\\R85r"SS\ \\\\45r\4S\ \S\\\ \\4\4/\\ \\4\\444Sjjr!S\R8S\\"\#\44Sjr$g)N) OrderedDict)wraps)CallableGenericOptionalProtocol) Concatenate ParamSpecTypeVar)_State)_get_root_modules_TT) covariant_PcHUS[[R"553$)N_)struuiduuid4)strings ^/var/www/auris/envauris/lib/python3.13/site-packages/torch/distributed/_composable/contract.pygenerate_state_keyrsXQs4::<() **c\rSrSrSrg) RegistryItemN)__name__ __module__ __qualname____firstlineno____static_attributes__rrrrrsrr_TStater )boundr_Mcx\rSrSrS\R S\R S\4SjrS\ RS\ 4Sjr Sr g ) _ContractFn&argskwargsreturncgNr)selfr)r*s r__call___ContractFn.__call__'s3rmodulecgr-r)r.r1s rstate_ContractFn.state)s3rrN)rrr r!rr)r*rr/nnModuler#r3r"rrrr'r'&s-FbggFFrF6BII6'6rr' state_clsr+c^[T5S[[[[4[4S[ [[[4[[ 44U4Sjj5nU$)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() funcr+c >^[T5S[S[RS[RS[4UU4Sjj5nS[ R S[4U4SjjnX!lU$)Nr1r)r*r+c>Un[U[R5(aU/nO[[ U55nT"5n[ 5n/n/n/n UGHun[ 5n [ 5n URR[U 5n [U [5(d[SU 35eURR[U 5n [U [5(d[SU 35eTU ;dTRU ;a[STRSU35eU RTU5 U RTRU5 UR[ UR555 UR[ UR!555 U R[ UR#555 GMx T"U/UQ70UD6nUcUn[U[R5(aU/nO[[ U55n/n/n/nUH{nUR[ UR555 UR[ UR!555 UR[ UR#555 M} [%U 5n[%U5nUU:wa[TRSUSUS35eS[[&S [[&S [&4S jn[)UU5H>unnU"[ UR+55[ UR+55S 5 M@ [)UU5H>unnU"[ UR+55[ UR+55S 5 M@ [)U U5H>unnU"[ UR+55[ UR+55S5 M@ U$)Nz-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 modules orig_fqnsnew_fqns check_keycX:Xag[U5[U5pCX4- nXC- n[U5(d[U5(a[USUSU35e[USUSU35e)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)r<r=r> 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: ) isinstancer5r6r listrr__dict__ setdefault STATE_KEYdictAssertionError REGISTRY_KEYrappendnamed_parameters named_buffers named_modulesrArzipkeys)r1r)r* inp_modulemodulesr3 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_modulesrGorig_named_paramsnew_named_paramsorig_named_buffersnew_named_buffersorig_named_modulesnew_named_modulesr9r7s rwrapper(contract..inner..wrapper_s  J&")),,!( ,DL9KE(NM DF !DF "AC "!URR[05RT5$r-)rKrLrMget)r1r9s r get_state*contract..inner..get_states+??--c$i r) rr%rr)r*r5r6r r3)r9rnrrr7s` rinnercontract..inner[sq tC C 77C iiC  C  C J bii F  " r)rrr r%rr'r#)r7rts` rcontractrv,sa^ 9Q{2r6*B./Q [R("g5 6QQf Lrr1c$[U[S5$)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)getattrrP)r1s r _get_registryrys 6< ..r)__composable_api_state_key)%r collectionsr functoolsrtypingrrrrtyping_extensionsr r r torchtorch.nnr5#torch.distributed._composable_stater torch.distributed.utilsr rrrrMrPrr#r6rJr%r'typervrNrryrrrrs$ #88== 65 TT"t_+  !#    )8t < T299d299o.7(GBGO47 &CG}C  k"b&!2%&' BF#R013CL/"))/c<6G1H(I/r