From 10965730437604f12d1f6e06ad349351ffda785b Mon Sep 17 00:00:00 2001 From: jdeast Date: Thu, 13 Aug 2026 22:45:42 -0400 Subject: [PATCH] Let the PyMC backend import without jax celerite2/pymc/ops.py imports pytensor.link.jax.dispatch at module scope, which imports jax. That makes jax a hard requirement of the PyMC backend at import time: >>> from celerite2.pymc import terms File "celerite2/pymc/ops.py", line 22, in from pytensor.link.jax.dispatch import jax_funcify ModuleNotFoundError: No module named 'jax' The only thing needed from that module is jax_funcify, and it is used solely to register the JAX conversion -- which is useful exactly when jax is installed. The conversion body already imports celerite2.jax.ops lazily inside the function, so this is the last eager jax reference on the PyMC path. Guarding the import and registering conditionally leaves the JAX behaviour identical where jax is present, and lets the PyMC and numpy backends work where it is not. That matters on platforms jaxlib does not ship wheels for -- macOS x86_64 stops at jaxlib 0.4.38, and jaxlib has no sdist -- where installing jax at all is not an option. Verified both ways against an installed celerite2 with this file patched in: with jax importable, _CeleriteOp is still in jax_funcify.registry; with sys.modules['jax'] poisoned to simulate its absence, 'from celerite2.pymc import terms' succeeds. --- .../pymc/__pycache__/ops.cpython-312.pyc | Bin 0 -> 13342 bytes python/celerite2/pymc/ops.py | 14 ++++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) create mode 100644 python/celerite2/pymc/__pycache__/ops.cpython-312.pyc diff --git a/python/celerite2/pymc/__pycache__/ops.cpython-312.pyc b/python/celerite2/pymc/__pycache__/ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb5de2c63b0ae0a3f4bbc17abd82d9aa66d180b8 GIT binary patch literal 13342 zcmc&)dvH|eb^q>b-%mYPPgdvw3wj{&6ku$DKnM)BYsWE4Jlkb=C01x3zIO#kUbV90 z))FN~O5z~Kq?S{6ECL;iHq*$RPD7~O$xLS2T`#HR8$5NJbf%f<9|9?faXU@V`R={@ zFiV5ej(es3?z!LdeBV9icg}ad`y;E>MBw`O@!uQ#n_5DCiVxbQQ;C%~Swb$50TLkJ zX8K5g39v&NC?ae#jEG_F2QWKAT$C z9(MFORN5IX>MLT15qlPfio+#+B~YS!iS(7Rq=y9bFOh&DXd7(Rn9Sam*wTx)BJMCcje_b(EJL>uguGDr1@2lU!9lV6sUo> zwJuolLIsQSDH(j6PA)GAy0mo)uG1)njIbkYJqeaJ5E12w?JGG(&i(o{Il~f)WL>|{ACHQ%MY(&Q8x95IvRR4_oeFw~qNjtR z_Anfa!Gk62i-(7Y)FSOkDJvNWMuMVm$eUN^)48$%5mYjKUild$7eG7`3@&GYrsLIdRbUAjRD}czGgpGu z7BAawHsoa9om^h%87ae6gT}i|BT{Kl#J9csSPPAMywWz3k+zP0fB*K*tsO1xzV;n0TQ+Uo+;Y6XJ|3u6eX!8e$f?_Bh>}XSZ%F77*+oCb4H8vvKyb9aisOXh~aj^>c ztB0Qyf@GX5>1`9wy)Zge_rh3GUwz9}`MTkvVYc)&`-FaqnKVtcpSAwXRk`SDPPv+I zS#0NfCwr#`W_;6wv(D-8>=)9OjoFAK131nGxJ%3g>1P7`OQt?PAb>6l6n}Myy@2Wl zNDUL@4AWh?_#qD~{{_jb42V3F zAYP!AE5X83M(7h9JY@tw!4B%Rdf7YL7fnRPp14sfNx)pFD)FO`6~}Sum7gDA;aR|X zaM!7KJ>IbukD(g@9 zB#@L-2gEf{BJ)Vgra&khj7XtqM3Qx(2uiwGj>XpCnYifls?zV#(MBSQmDq?$6DC+w zTnCBAsCFvFf`0lYWNrXN5ez+y6bSj_s9UGJQP4hV7{Sos^@bv$xYr9x$xFeZer165 z*rOX;8NFUmElKivB>|p6;oRe5B^I|pGFp-&46RyUo1jE`6OxZX>`Pn|{5hj`E3Jel z(?@3e`5luxrg~?FZ<-r#RW>d;N~WHkw#=MLIT~*{N|q{WX8JEy&*@SXtxMH)OKa+9 zA5E=kU#eO&Yq+>)sjPCQb2@sPW2&1z0{Y=p*~Z%fTip!0P+mog^jvk*vVqjp&AKig zU0U0)T5NI_tq9gC+eG)d!>hH#Xq(ih46Yv=D#dDG)SnSRc(B3`#cm}4VtavPu`|JF z=$n&Q5T_M1*J_^x_d?tq$MvIEf(&K^M$04QH86mmtzT%t<+R9f_Zs<@ppE?^lT{E< zaZtYQj7VyAMArlGCVVSnd!Wy(a)B}zvQ+8{sH)seo;^aym*HuQ9pgZ)E!lZ%RL)E| zMKLx0aQ!6a+}bRc0eyl~6cyIMc>k7AYAQS$tdUOZ|3!uet5t9oW&{Ip?Ksu7Yk(PW z+SmYOoZ{pKtXL0!gBbwO<^`al)uK)w)&iHXi6>(`^i_CVJJz~SnFGrZtdOtsf}a~e zr?VY7V@ccz1@sy$(x~h`-~?VMkU?9)&(fA4>tjNKPf$>tX3rU^vnputrU4j~YzCfn zW89b!FV*@SEYqm$9X)`APIFkF1n#IK$DP(FL8YLlG06249i>1vYH|n90CL5H0r#lY zt>9p@dmtKj=P4*!lg;p3*$K&E@CVV2dbSv*!6cZ$EFT#vKv0Cf$y{K-cP!<}2=iy$ zX~x5f$ox^!0qJE{W?MJOoKGB(P-xw5igL1XH@bvpV&b0B^|@FV_MHrRBhf&xb@x!z z?;Dc#v}S68s1pe>^4|~vleQh#%I4}Pr1K+_BQw1}wl`iYyDOoS_F_d9Gg{>qgJ3+t zcW|rjjyom>{UJ04B7M{Ky%B#@guJ*rG8{f06y4E&ch+FI`=g?p!uF_nSIz?Jl$E1) zLe}}CQ86GBvZZs4FYYzj< z*5hK-7x06BwdfUa_{_!RP)2>^^YXibuANSFH?T2ymDsS*PovX zFLR8uafRcGZOeqS+Exw3Vo%}gOWStbH1B}!X8E^_Zy0BuPj1+stT~V_>PcF9ZdFvz zJU#o=TgTowHg_h~xZ@uyIwyN4_Dz*6*^ADfnLKl8^GtlU=i*qpswHi2xmD&~D(hTm zSUB}g^Fl2!aFB7|a@Wsx&y~M%B<=1)%J+V5c$|fJ3X__sWJ@H2EJd>*HNL6lL ztn5lvcA=GhX4*FMT*^^*n`5hNw+U;tfz_VsJ2Md~WagBa=s#92N6{h3$VHPFHnb3(q@ej?C4itDEP;sjBW3 z4yyeE26Es)hZr2mew#o?s}|xYJ^%dV^D|p!%jU|`_DxB{rn`44fQh?z9He?bb9Y66 zdJ?)X?+Ts+Yx(zzY7R8(-fOcQSXc7iE-TP~w@!dG4<1mM@VD|2Bro&ma4Obl3|!5E z@W7&;<)A$1E=76jPJ)BV1$eBXxQ$3sQR%0Hj=v0Q@u`%}g@jO~KL|T6+!GeS;w&nkgq9Kl{>KU|JGWo1nJt-W zn2Nu;;d0I06-4B(*oznnx^!b7AC3(LJt{*Kw+>?;IJt}$O!?XH;QXNcN!XXd1!#4t z>7QM_Q~WANtVL<(_PLU|f$!7+v1r-;fvZ=s6kn-pV6un}hKq88cE%%6fR%YPT(VBt z8Y{(=Sa&T5U8--qo0@;Hq=g8%L;ojN}4n$V&4=Ae$iIe_mrfC1pUAm+?> zFSxE9z22VMvOiUMAZ0nQZSHAS*}v1UbpVyNRV!5SF>5sX06 zJ;=M^O>5B?kktN63}94}P|{x*BB@?>jF=qI>-77-$jh()jaaQ;A8p&`V*VZ7wXO z2#yd)}(DC88vXPK=0Wow6WATrMB3(lR%I?AqJ$ekP1;gD`PC$z; z&{3!$^B6}LsW~&iGc42Z!cVnmdR6d8iwnHOC{&l~VV7`soP1^jwpiJjszguy%sZ_= zbf$J5PE|gdvOM}JGX_31)hF z*QC~Ur5s%l$gQ@m7@^_{(OXTo1tXfq*y7w*B%X+Bri3He2zTw%siN9g5+`E}J-wDP({rImJ)bCLduEtcs6MXg)K=HjiF5LA)RW z+wq7$DD$|dmJM)NLAT_xVLP#8g2LrOQ+=+oVgCi}&N*l%Jq1Yymz<{C7Tg0)^8ifk08AydGrrgRFZREA zVzFs^s%iUT)4^2J!JEas;0-&AFIcCoGXt}}i-R{Eo{vh)Rs_~rj$7SB%v}lJ{8ze4 zyIcxlfgk(@$Pss7LZM?LQn3(y)^TqCF-I$>WiV*DIU6-f%Q9JcP8x3kRFNA+-de z@`$t$m8W5NaUah0aZH}YBm{|URU`Lsd^H5|eMW49TnW`sK~8!qK$8|MXn^G7m&lT% zZoK!S=JsEbQlaWjYn9M`XQxpx-6^XO8h*LWB{` zl21hjEPom?FYe6q?sVl&7&md92QMe@p%ySL^ zGVZi5fO;@e{$2=swi-jN2a_j>Lmv%qPaBRvMr<>c{u1 zUHJBj;kU9E5}#68sgl$0vybjjU`?1rUOl$3u~TM<+9J zneoevPiDR#Gf&HUJgNt7DWm|-WsvSU7-iJs%{j^UYV*v(T~@s3p}Mpd65KZ1%NFfx zQuZ~A_J)+bVfKkRX0CnCH}6c_Hzy67l|v`=sI?0>V07~6I?~CTKKtae5^+WCJG%_p z=Mi}R7Zaa{QNXQBz)c*(r?jibI!T0x##BV=AT$#7TOnl*Nd+IPTGDjX*j4X z+vgxN)?9Q_R_n4ZXSL`;{Q+<a`_lFuNy83hB{E734e^Cj0Xhf-6(GZJD?nG&ezU^k%;hvIL6?n% z)pm_=waS`m~ft!2JVGJ7cH*8rq1jU{Vc<|MQjs zg^#QZh!`l8?*Do-^*_@*q5tpW6u+67=jWjHwmd)Meri9CETM=OR(y=MjJEv-#m^fE zH2(%P%G3PDl)Z7$z9D7bpqLF5=Cr*tY3Tf%-+ECVtJ&9nuj-HElF)Ve8d6`!z$K>&g?IcToaLLDsUm`E7eXI3 zdEG!zlc)P_OrG{VSJvE>72HCofS{Mzcsu+$6vrv!Gv#0cda{bKQ}@gmq8r`*8$hAp zpa2ek4~y|7Df(Wc?~IopY?OuEXiyVjVI< zpMu8X+nAuxdn_XQ=i($L=xvFwVuBV`d<_!>V)0KfK@}4(V}kENq)gNujjEfH2t+0E zV+;5WV<05K1_i!8!A~k+qu|us{f<`H1hm2)LFRyvIiQp%r{QpLgoZNZ;-k@cZx{{? z;X^+7fQ4pyoFeWa#VrS4=GrIo`*RnqBi^W2b0R>Z?(5)h7+4K?RoUgQ?tcmC{ zeFuIL5+9F~Um1w7KFP0JG?Fd6LSyJ+#FJ=*2BQ(cI9<4%*7C@6il~lbxXFk zdG>qyYx=p7q@@!=t9CdAWR0d}4$41~EWrH_%57aun&`0T4@$BzOY$^2VTA8_#dtJ2B+(F&EI>RvHliH69EaXzyB0ocRikJ+ zPvp`6S;kU=4 z4J(*Is9-)yyd2KyiOfD0Q{pc|e`p9MXBdpar&7N737L<^f)IJ>mwEVbIv|^Kf-*Gr zBI4GWMp{I4r$mGa5y47CC>0UEsffr1O&1_+v>tM4*l3d?T87<9P->6(4^RS;5h6`N z3Sno4xuY|h*>U|Z$_TUmC&c_Asro4?`U$arNUHydH2f=R{E*Zw^E#&dLw?&Ghn}fe zEg4{#T~li>G)_0J5~P=(;)t$v{MgIh7rawnN(*(1!iJQvA-VCPw6JGU*q0LaUEg<8 zIQYY^rH)-{sIBX}&FP|tuQks#Ol_IbP4Av9NfkHFH>Zjo9zXmm=&qN6sG+HJ>dUjc z;0vAl9cg3dxb6dc<;?E6s`+P9o;}wCH|>YVEg#q_l9e0hzO+!2YUxSa4vw1(BJWD^u9>PdJU8w7F~1uQ+WD%Fm*HDu1Lz6Ql!L|;D_)zQlk4 literal 0 HcmV?d00001 diff --git a/python/celerite2/pymc/ops.py b/python/celerite2/pymc/ops.py index cce6928..da45f64 100644 --- a/python/celerite2/pymc/ops.py +++ b/python/celerite2/pymc/ops.py @@ -19,11 +19,18 @@ import pytensor import pytensor.tensor as pt from pytensor.graph import basic, op -from pytensor.link.jax.dispatch import jax_funcify import celerite2.backprop as backprop import celerite2.driver as driver +try: + from pytensor.link.jax.dispatch import jax_funcify +except ImportError: # pragma: no cover - jax is an optional dependency + # pytensor.link.jax.dispatch imports jax, which is not a required + # dependency of the PyMC backend. The conversion defined below is then + # simply not registered. + jax_funcify = None + def _resize_or_set(outputs, n, shape): if outputs[n][0] is None: @@ -160,7 +167,6 @@ def grad(self, inputs, gradients): # JAX conversion for PyTensor JAX linker ------------------------------------- -@jax_funcify.register(_CeleriteOp) def _jax_funcify_celerite(op, node, **kwargs): """Map celerite2 PyTensor ops to their JAX counterparts.""" @@ -224,3 +230,7 @@ def general_matmul_upper_fwd(t1, t2, c, U, V, Y): raise NotImplementedError( f"No JAX conversion registered for {op.name}" ) + + +if jax_funcify is not None: + jax_funcify.register(_CeleriteOp)(_jax_funcify_celerite)