VpfEVUddlmZddlZddlmZmZmZddlZddlZddl m Z ddl Z ddl Z ddl mZmZddlZddlZddlZddlZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZdd lmZddlmZddlm Z ddl!m"Z"ddlm#Z#ddlm$Z$ddlm%Z%ddlm&Z&ddl'm(Z)ddl*m+Z+ddl,m-Z-ddl.m/Z/ddl0m1Z1m2Z2m3Z3m4Z4m5Z5ddl6m7Z7m8Z8dZ9dZ:dZ;ej<e=e)j>Z?e?j@ZAe)jBZBe)jCZCejDeEcZEZFejGeHcZHZIejJeKZLdZMdZNejOdud"ZPd#ZQeZRGd$d%ejSZTeTZUd%eVd&<ejWd'ZXejYdvdwd-ZZdxd3Z[dyd7Z\e]Z^d8eVd9<dzd:Z_Gd;dZad|d?Zbd}dBZcdZddCeVdD<d~dEZeddHZfddIZgddLZhddOZiddPZjddUZkddZZld[Zmd\Znejod]^Gd_d`ZpdaZqddeZrddiZsejtdjZud]eu_veuweseuxdkdlZye ejzeuej{eu<eyej|eu<dmZ}e}e j~eu<dnZe#jeuedope#jeuedqpdrZe#jeuedddsdtZee%jeu<dS)) annotationsN)CallableIteratorSequence)partial)Any NamedTuple) basearray)config)core)api)dtypes)source_info_util)traceback_util)util)ad)batching) array_types)mlir)xla)pxla)lib) xla_client)record_event_duration_secs) PartitionSpec)Sharding)SingleDeviceSharding NamedSharding GSPMDShardingTransferToMemoryKindis_single_device_sharding)LayoutDeviceLocalLayoutz&/jax/core/compile/jaxpr_trace_durationz//jax/core/compile/jaxpr_to_mlir_module_durationz*/jax/core/compile/backend_compile_durationFct|fi|}tjd} ||}tj|n$#tj|wxYw|S)zEImpl rule that compiles and runs a single primitive 'prim' using XLA.F)xla_primitive_callablerjax_jit#swap_thread_local_state_disable_jit)primargsparamsfunprevoutss Q/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/jax/_src/dispatch.pyapply_primitiver/Psxt..v..#  8 8 ? ?$: 3:DK33D9999CK33D9999 +s A!A4r(core.Primitivec jfd}j|_j|_tj|S)Ncj|iSN)bind)r)r*r(s r.prim_funz(xla_primitive_callable..prim_fun^s 49d %f % %%)name__name__ __qualname__r jit)r(r*r5s`` r.r%r%\sF&&&&&&i()(   r6cV|tt|dSr3)def_implrr/)r(s r. simple_implr=es$--../////r6cReZdZUdZded<ded<dZdd ZddZddZdZ dZ dS)RuntimeTokenSetzJSee docstring for effects.py module for the calling convention for tokens.zdict[core.Effect, core.Token]current_tokenszdict[Device, RuntimeToken]output_runtime_tokensc"i|_i|_dSr3r@rAselfs r.__init__zRuntimeTokenSet.__init__uD!#Dr6eff core.Effectdevices list[Device]return core.Tokenc|j|tjdtj}t |t jr2tj |tj |Stj j |}t jtj|g|gd}||j|<|S)Nr)r@getnpzerosbool_ isinstancer Tokenjax device_putshardingPositionalShardingrget_replicatedr shard_args)rErHrJtoks sharded_toks r.get_token_inputzRuntimeTokenSet.get_token_inputys   ! !#rx28'<'< = =C#tz""K ^C!@!@!I!I J JJ "11'::A*T_aS3%88;<.''h'11 TT r6c$eZdZUded<ded<dS) SourceInfozsource_info_util.SourceInfo source_inforteqn_nameN)r8rlr9rnror6r.rrs%****-----r6r%Iterator[tuple[Sharding, SourceInfo]]c#pKddlm}ddlm}|jD]oj|jur1tjjj j dfVBj|j ur^tjjj fdj dDEd{Vfdj dDEd{Vj|j urgj d j stjjj d fd gj d j d DEd{V#jtur>tjjj fdj dDEd{Vqtj|D]}t#|Ed{VdS)Nr)pjit) shard_maprWc3 K|]}|fV dSr3ro).0irs r. z-get_intermediate_shardings..s(GGq1k"GGGGGGr6 in_shardingsc3 K|]}|fV dSr3ro)rors r.rz-get_intermediate_shardings..s(HHq1k"HHHHHHr6 out_shardingsmeshctrtdznd}tfdt|DS)Nrc3BK|]}|VdSr3)rO)rrnamess r.rzFget_intermediate_shardings.._names_to_pspec..s-BBuyy||BBBBBBr6)maxrrange)rndmins` r._names_to_pspecz3get_intermediate_shardings.._names_to_pspecsA"'.E QQBBBBU5\\BBBCCr6c3dK|]*}tjd|fV+dS)rN)rr*)rrrrrs r.rz-get_intermediate_shardings..s_TT!F!3__U5K5KLLkZTTTTTTr6in_names out_namesc3XK|]$}t|tr |j|fV%dSr3)rSr memory_kind)rr\rs r.rz-get_intermediate_shardings..sRLLq8,,L121Jk"1J1J1J1JLLr6rJ)jax._srcrjax.experimentalrrrsharding_constraint_prrr7r*pjit_p shard_map_p_is_jax_device_mesh device_put_pr rget_intermediate_shardings)rrrrrrrs @@@r.rrs(((((( ZLLc }222s 0BCCk Z #[ 11111 $+ % %s 0BCCkGGGGCJ~,FGGGGGGGGGGHHHHCJ,GHHHHHHHHHHH )/ / / Z  3s 0BCCkDDDTTTTTTRJ!7R#*[:QRTTTTTTTTTTT , & &s 0BCCkLLLLCJy,ALLLLLLLLLL.''44h)(333333333344r6c td|jDp>tdtj|gt j|DS)Nc3K|]H}t|jtj!t |jjtjuVIdSr3)rSavalr UnshapedArraytypedtypebint)rvs r.rz"jaxpr_has_bints..s\99!AFD$6779d16<  DI-999999r6c3K|]U}|jD]K}|jD]A}t|jtj!|jjD]}t|VBLVdSr3)routvarsrSrr DShapedArrayshape_is_bint_axis_size)rjerds r.rz"jaxpr_has_bints..sNNvNNqyNN"#AFD$566NAB NN<=!##NNNNNNNNNr6)anyinvars itertoolschainr rrs r.jaxpr_has_bintsrs 99u|999 9 9 N NN wu0E0EFFNNN N NOr6r core.AxisSizecHt|tjr)|jrJt |jtjuSt|tjrDt|jtj o$t |jjtjuSdS)NF) rSr DArrayrrrrVarrr)rs r.rrs4;-w ==DI %%!TX- qvt0 1 1 ,   $) +- r6z)Callable[[core.Jaxpr], core.Jaxpr] | Noneoutfeed_rewriterc2tt|S|Sr3)rrs r.apply_outfeed_rewriterrs! E " "" Lr6argrct|tjs7tj|s%t d|dt |ddSdS)N Argument ' ' of type z is not a valid JAX type.)rSr Tracer valid_jaxtype TypeErrorr)rs r. check_argrsv S$+ & &!$* > >>r6r core.JaxprEqnc|jd}|r+|jddt|zS|jtjvrt |jSdS)N call_jaxpr axis_sizer)r*rOrrrinitial_style_primitives!_initial_style_primitive_replicas)rrs r.rr*sez~~l++* :>>+q ) )N:,F,F FF }444 ,SZ 8 88 1r6r*dict[str, Any]cxttjt|dS)Nrr)rr traverse_jaxpr_paramsrrj)r*s r.rr3s7 T ' ? ? F F H H   r6cFtjjptjjSr3)r debug_infsr| debug_nansror6r.needs_check_specialr7s   ;F$5$;;r6r7bufsSequence[basearray.Array]Nonec\tr|D]}t||j|dSdSr3)r_check_specialr)r7rbufs r. check_specialr:sJ+++T39c****++++r6rnp.dtyperbasearray.Arrayctj|tjrtjjrJtjtjtj |rtd|tj jrLtjtj tj |rtd|dSdSdS)Nz#invalid value (nan) encountered in z#invalid value (inf) encountered in ) r issubdtyperPinexactr rr|risnanasarrayFloatingPointErrorrisinf)r7rrs r.rr?s ubj))M M26"(2:c??*C*C#D#DM KTKK L LL M26"(2:c??*C*C#D#DM KTKK L LL MMMMMMr6c|Sr3ro)xs r. _identity_fnr Gs (r6c ddlm}m}||j}|j|jkr |jt||S|j|jkrd|jD}|jdj }d|jD}|jdj }td|d|d|d|| |j } | r| } n tj|jjt$g |j} | } g| _g| _tj| | | _t0j| } t7tj|j| t7tj|j| jksJ||jt=|j| |j |j } |jt|| S) Nr)r arrayrcg|] }|j Sroidrrs r. z3_different_device_order_reshard..Ts===qt===r6cg|] }|j Srorrs r.rz3_different_device_order_reshard..VsCCC1!$CCCr6z[Input and target sharding should have the same set of devices. Got input's device set ids: z on platform z' and target sharding's device set ids: )otypes)r)!rr r_check_if_deletedrW_device_assignmentr:r  device_setrupper ValueError_to_xla_hlo_shardingndim is_replicatedrP vectorizeindexrto_protoiota_reshape_dimsiota_transpose_permtaketile_assignment_devicesxc HloSharding from_protolist$make_array_from_single_device_arraysrrr_arrays)r target_shardingr r inp_shardinginp_idsinp_plat target_ids target_platold_hlo_shardingnew_hlo_sharding permute_ordernew_op_shardingnew_xs r._different_device_order_reshardr7Js!!!!!!!!,$(JJJ ?737< ? ? ? B BB :::==\<===G.q1:@@BBHCC BCCCJ!4Q7@FFHHK D=DDD!)DD(DD7BDD E EE "66qv>>##%%G'/BL!C!I),////;/NPPM'//11O(*O%*,O'.0g'??AA//O+~00AA 8)AACCEE F FBGO>+CEEFF F F F F  4 4gO68H / ;===i   % >_ = = =e D DDr6T)frozencReZdZUdZded<ded<ded<ded <ed Zd S) _DeferredShardArgzDeferred call to `pxla.shard_args`. Per-array impls return this object instead of a result array to indicate a deferred `shard_args` call. `_batched_device_put_impl` then batches all `_DeferredShardArg` objects into a single `shard_args` call. rr rr\zcore.AbstractValuerr committedcLtj|j|j|jSr3)rglobal_aval_to_result_handlerrr\r;rDs r.result_handlerz _DeferredShardArg.result_handlers  -di P PPr6N)r8rlr9rmrnpropertyr>ror6r.r:r:sg &&& +++/// QQ 8QQQr6r:c$ddlm}ddlm}t |t r|}t |dd|krt |ddr|S|jsCt ||jr.|js't |t sJt||S|jrt ||jrp|jrit|j dkrQ|j |j j kr<|j |j j kr't |t sJt||S|jst ||jr|jrt|t vr[||t|dt|d  t%jt(| |St+d |d t-|||dSt ||jri|js$t+d|||St1|j r&t3j|t7||g|gSt7|t3jn|}t-||||duS)Nr)r)multihost_utilsrW _committedFrzc passed to device_put is not the same on each process. Make sure you are passing the same value of z on each process.) fail_messagerzjdevice_put's second argument must be a Device or a Sharding which represents addressable devices, but got zF. Please pass device or Sharding which represents addressable devices.TzZdevice_put's first argument must be a fully addressable array, but got value with devices )rrrrArSrgetattris_fully_addressable ArrayImplr7lenr_internal_device_listrWrBrr assert_equalr r:r rr:rJr!rbatched_device_putr_get_default_device)r rrbrrAr\shs r._device_put_sharding_implrMs......!! /Aq*d##q((WQ e-L-L( h "31eo&&3/0/E3 8 $ $$$$ ,Q 2 22 3:a#A#A3 3#&q|#4#4q#8#8 1:#CCC  --- 8 $ $$$$ ,Q 2 22 != a ) )9!,9 q''[ $$ 77//GG/// % 1 1 1 6sw|1555a888  <67 < < < = == Q4 . ..5?## / !3  2$%IIKK 2 2 3 33~ h "1: . ./  $T+?+G+G!&,X///%~!466639;;" 1b$d(: ; ;;r6rb!Device | Sharding | Layout | Nonesrcc |t|tst|trtd tj|}n6#t $r)}t d|dt |d|d}~wwxYwt|tr|}|j}t|dr |j jnd}||j t|||j St|j tr$t|tt dfs'td|d|t!|dd|krt!|dd r|S||t|||j St#jt&| |St|||S) NzTransferToMemoryKind argument to jax.device_put can only be used inside jax.jit. If you are using device_put outside jax.jit, then please provide a concrete Sharding with memory_kind.rrz is not a valid JAX typelayoutzVsharding and device_local_layout in `Layout` instance should be concrete. Got layout: z for input rBFr)rSr rr abstractifyrrr"device_local_layouthasattrrQrWrMrr# str_shortrDr r:r )r rbrOrerrldllx_dlls r._device_put_implrZs -..A*++A  @ A AA N ?1  DD NNN CQCC$q''CCC E EJMNN5A C,3Ax,@,@ JAH ( (dE {qz) &q$ ; ;; qz8 , ,F s.T ; < <F  E$% E E26..2B2B E E F FFq(D!!Q&&71lE+J+J& h } &q$ ; ;; 137>BrJ+Sequence[Device | Sharding | Layout | None]srcsc@g}ggg}}}tt|||D]\}\}} } t|| | } t| trI|||| j|| j|| |rftj ||} t|| D]@\}} t||tsJ|| | ||<A|S)N)rbrO) enumerateziprZrSr:appendr r\rrZr>)rJr\xsysshard_arg_indices shard_arg_xsshard_arg_shardingsrr rbrOyshard_arg_resultsshard_arg_results r._batched_device_put_implris;  "9;R#6\&s2w'='=>>a !VS6s333A!&''&q!!!!#  %%%IIaLLLL5(;\JJ"#46GHH55  10 1 1111e""#344bee )r6rVc|Sr3ro)rJr\ras r.rk s"r6c|dgt|z}g}tt|||D]=\}\}}} t|tjur||||| f>|rHtt|\} } }}tj | ||d} t| | D] \}} | ||< |S)NrJr\) rGr^r_rrZeror`r)rr4)ctsrJr\_resultsdp_argsrctrbrOindicesr)rbrfs r._device_put_transposerus FSXX ' ''C$(?(?@@++a "fc Bxxrw nnaVS)*** #'W #6#6 GT7D  D$W = = =BGR  1gajj .r6c d|Dr-tfdddDs J|tj|i||fS)Nc.g|]}|tju|Sro)r not_mapped)rbds r.rz'_device_put_batcher..s%PPPb"H. s?&&%'b &&&&&&r6r)allrr4) batched_args batch_dimsr*r{s @r._device_put_batcherrsPPJPPP #&&&&+.lower's6H&:;<<& FH % %Q  & Hf99$)DDMMOO Q Q  $Q(:H E Ea h Hr6)r)mapavals_in avals_out)rrJr\rars` r._tpu_gpu_device_put_loweringr&sB      c%WdCL#-HH I IIr6r)rgpuc|D]A}t|ttfr#|jt d|jjB|S)NzNPassing memory_kind to device_put via Shardings is not supported on platforms )rSrr rNotImplementedErrormodule_context platforms)rrJr\rarbs r._common_device_put_loweringr7sd88f6H&:;<<8&  7*4 7 7 8 88 )r6rmcg}|D]N}t|ttfr||j9|dO|Sr3)rSrr r`r)rJr\xm memory_kindsrbs r._propagate_mem_kind_dprAse,  f&8%9:;; &,----$ r6)r(r0r3)rsrtrurtrvrw)rrrrtrLr)rrrrtrLr)rr)rrrLr)rrrLr)rrrLr)rrrLr)rr)rrrLr)rrrLr)r*rrLr)rLr)r7rtrrrLr)r7rtrrrrrLr)rbrNrOrN)rJr[r\r[) __future__ratexitcollections.abcrrr contextlib dataclasses functoolsrrrtypingrr r} threadingnumpyrPrUrr r r r rrrrjax._src.interpretersrrjax._src.abstract_arraysrrrrr jax._src.librr&jax._src.monitoringrjax._src.partition_specrjax._src.shardingrjax._src.sharding_implsrrrr r!jax._src.layoutr"r#JAXPR_TRACE_EVENTJAXPR_TO_MLIR_MODULE_EVENTBACKEND_COMPILE_EVENTregister_exclusion__file___xlaxeClientBackendrcCompileOptionssafe_maprrsafe_zipr_ unsafe_zip getLoggerr8rrzr/cacher%r=rdlocalr?rprnregisterrrcontextmanagerrrrsetrrrrrrrrrrrrrrrr r7 dataclassr:rMrZri Primitivermultiple_resultsr<def_abstract_evalru linear_jvpprimitive_jvpsprimitive_transposesrprimitive_batchersrregister_loweringrrmemory_kind_propagate_ruleror6r.rsT #"""""" 8888888888 """""""" %%%%%%######$$$$$$******000000&&&&&&%%%%%%&&&&&&))))))::::::111111&&&&&&DDDDDDDDDDDDDD65555555=ND!!(+++W ) "-Z-Z  8 $ $     000 44444io444l#2/"3"33333%%%  6 6 6 6 6>ASUU%BBBB 4444>OOOO?CBBBB!!!! ????    <<<<++++ MMMM   2E2E2Ejd###QQQQQQQ$#Q$4<4<4>>>,u>>>>   |%@AAA(,41G ---r6