Coverage for src / bartz / jaxext / __init__.py: 91%

1# bartz/src/bartz/jaxext/__init__.py 

2# 

3# Copyright (c) 2024-2026, The Bartz Contributors 

4# 

5# This file is part of bartz. 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the "Software"), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in all 

15# copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

23# SOFTWARE. 

24 

25"""Additions to jax.""" 

26 

27import math (empty)

28from collections.abc import Sequence (empty)

29from contextlib import nullcontext (empty)

30from functools import partial (empty)

31 

32import jax (empty)

33from jax import ( (empty)

34 Device, 

35 debug_key_reuse, 

36 device_count, 

37 ensure_compile_time_eval, 

38 jit, 

39 random, 

40 vmap, 

41) 

42from jax import numpy as jnp (empty)

43from jax.dtypes import prng_key (empty)

44from jax.lax import scan (empty)

45from jax.scipy.special import ndtr (empty)

46from jaxtyping import Array, Bool, Float32, Key, Scalar, Shaped (empty)

47 

48from bartz.jaxext._autobatch import autobatch # noqa: F401 (empty)

49from bartz.jaxext.scipy.special import ndtri (empty)

50 

51 

52def vmap_nodoc(fun, *args, **kw): (empty)

53 """ 

54 Acts like `jax.vmap` but preserves the docstring of the function unchanged. 

55 

56 This is useful if the docstring already takes into account that the 

57 arguments have additional axes due to vmap. 

58 """ 

59 doc = fun.__doc__ (empty)

60 fun = jax.vmap(fun, *args, **kw) (empty)

61 fun.__doc__ = doc (empty)

62 return fun (empty)

63 

64 

65def minimal_unsigned_dtype(value): (empty)

66 """Return the smallest unsigned integer dtype that can represent `value`.""" 

67 if value < 2**8: 175 ctx2J d e Z K 2 sctcuc;s}s~s=sk ata bto vkwk?sxkyk@sL b l 0 1 dbM f N O zk3 P Ak4 Q h R S i 5 BkCk[sDkEk]sFkVcGkHkIk^sw T U xcJk_sp j q yczcWcV c m W X KkLk`sv Mkrbsbtb6 ub8ss Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pcct9sdt!set#sft$s%s{s'sy + , z A hbB C D E F H ibx gt(s)s*s+s,s-s.s/s:shtitjtktjbkblbmbI nbobpbwbltxbmtybzbntAbBbotpt

68 return jnp.uint8 172 ctx2J d e Z K 2 sctcuc}s~sk ata bto vkwk?sxkyk@sL b l 0 1 dbM f N O zk3 P Ak4 Q h R S i 5 BkCk[sDkEk]sFkVcGkHkIk^sw T U xcJk_sp j q yczcWcV c m W X KkLk`sv Mkrbsbtb6 ub8ss Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pcct9sdt!set#sft$s%s'sy + , z A hbB C D E F H ibx gt(s)s*s+s,s-s.s/s:shtitjtktjbkblbmbI nbobpbwbltxbmtybzbntAbBbotpt

69 if value < 2**16: 69 ↛ 71line 69 didn't jump to line 71 because the condition on line 69 was always true75 ctx2J d e Z K sctcuc;s=sk a o vkwkxkykL b 0 1 M f O zkP AkQ h R S i BkCkDkEkFkVcHkIkw T U Jkp j q yczcV c W X KkLkv Mk6 8s9s!s#s$s%s{s's(s)s*s+s,s-s.s/s:s

70 return jnp.uint16 75 ctx2J d e Z K sctcuc;s=sk a o vkwkxkykL b 0 1 M f O zkP AkQ h R S i BkCkDkEkFkVcHkIkw T U Jkp j q yczcV c W X KkLkv Mk6 8s9s!s#s$s%s{s's(s)s*s+s,s-s.s/s:s

71 if value < 2**32: 

72 return jnp.uint32 

73 return jnp.uint64 

74 

75 

76@partial(jax.jit, static_argnums=(1,)) (empty)

77def unique( (empty)

78 x: Shaped[Array, ' _'], size: int, fill_value: Scalar 

79) -> tuple[Shaped[Array, ' {size}'], int]: 

80 """ 

81 Restricted version of `jax.numpy.unique` that uses less memory. 

82 

83 Parameters 

84 ---------- 

85 x 

86 The input array. 

87 size 

88 The length of the output. 

89 fill_value 

90 The value to fill the output with if `size` is greater than the number 

91 of unique values in `x`. 

92 

93 Returns 

94 ------- 

95 out : Shaped[Array, '{size}'] 

96 The unique values in `x`, sorted, and right-padded with `fill_value`. 

97 actual_length : int 

98 The number of used values in `out`. 

99 """ 

100 if x.size == 0: 16 ctx2e a b l c m qt|svcwcwbxbybzbAbBb

101 return jnp.full(size, fill_value, x.dtype), 0 1 ctx2qt

102 if size == 0: 15 ctx2e a b l c m |svcwcwbxbybzbAbBb

103 return jnp.empty(0, x.dtype), 0 1 ctx2|s

104 x = jnp.sort(x) 14 ctx2e a b l c m vcwcwbxbybzbAbBb

105 

106 def loop(carry, x): 14 ctx2e a b l c m vcwcwbxbybzbAbBb

107 i_out, last, out = carry 14 ctx2e a b l c m vcwcwbxbybzbAbBb

108 i_out = jnp.where(x == last, i_out, i_out + 1) 14 ctx2e a b l c m vcwcwbxbybzbAbBb

109 out = out.at[i_out].set(x) 14 ctx2e a b l c m vcwcwbxbybzbAbBb

110 return (i_out, x, out), None 14 ctx2e a b l c m vcwcwbxbybzbAbBb

111 

112 carry = 0, x[0], jnp.full(size, fill_value, x.dtype) 14 ctx2e a b l c m vcwcwbxbybzbAbBb

113 (actual_length, _, out), _ = scan(loop, carry, x[:size]) 14 ctx2e a b l c m vcwcwbxbybzbAbBb

114 return out, actual_length + 1 14 ctx2e a b l c m vcwcwbxbybzbAbBb

115 

116 

117class split: (empty)

118 """ 

119 Split a key into `num` keys. 

120 

121 Parameters 

122 ---------- 

123 key 

124 The key to split. 

125 num 

126 The number of keys to split into. 

127 

128 Notes 

129 ----- 

130 Unlike `jax.random.split`, this class supports a vector of keys as input. In 

131 this case, it behaves as if everything had been vmapped over, so `keys.pop` 

132 has an additional initial output dimension equal to the number of input 

133 keys, and the deterministic dependency respects this axis. 

134 """ 

135 

136 _keys: tuple[Key[Array, '*batch'], ...] 

137 _num_used: int 

138 

139 def __init__(self, key: Key[Array, '*batch'], num: int = 2): (empty)

140 if key.ndim: 934 ctx2J - d Gbe qbZ . K Hb2 IbTcNkOkk a o / JbKb: LbMbL ; b Nbl ObPbPk0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw QkT RkU SkTkUk} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW VkX Wkcb+b,bv Xkrbsbtb6 ubYkZk0k1k2k3k4k5k6k7k8k9k!k#k$k%k'k(k)k*k+k,k-k.k/k:k;k=k?k@k[k]k^k_k`k{k|k}k~kalblcldlelflglhliljlklllmlnlolplqlrlsltlulvlwlxlylzlAlBlClDlElFlGlHlIlJlKlLlMlNlOlPlQlRlSlTlUlVlWlXlYlZl0l1l2l3l4l5l6l7l8l9l!l#l$l%l'l(l)l*l+l,l-l.l/l:l;l=l?l@l[l]l^l_l`l{l|l}l~lambmcmdmemfmgmhmimjmkmlmmmnmompmqmrmsmtmumvmwmxmymzmAmBmCmDmEmFmGmHmImJmKmLmMmNmOmPmQmRmSmTmUmVmWmXmYmZm0m1m2m3m4m5m6m7m8m9m!m#m$m%m'm(m)m*m+m,m-m.m/m:m;m=m?m@m[m]m^m_m`m{m|m}m~manbncndnenfngnhninjnknlnmnnnonpnqnrnsntnunvnwnxnynznAnBnCnDnEnFnGnHnInJnKnLnMnNnOnPnQnRnSnTnUnVnWnXnYnZn0n1n2n3n4n5n6n7n8n9n!n#n$n%n'n(n)n*n+n,n-n.n/n:n;n=n?n@n[n]n^n_n`n{n|n}n~naobocodoeofogohoiojokolomonooopoqorosotouovowoxoyozoAoBoCoDoEoFoGoHoIoJoKoLoMoNoOoPoQoRoSoToUoVoWoXoYoZo0o1o2o3o4o5o6o7o8o9o!o#o$o%o'o(o)o*o+o,o-o.o/o:o;o=o?o@o[o]o^o_o`o{o|o}o~oapbpcpdpepfpgphpipjpkplpmpnpopppqprpsptpupvpwpxpypzpApBpCpDpEpFpGpHpIpJpKpLpMpNpOpPpQpRpSpTpUpVpWpXpYpZp0p1p2p3p4p5p6p7p8p9p!p#p$p%p'p(p)p*p+p,p-p.p/p:p;p=p?p@p[p]p^p_p`p{p|p}p~paqbqcqdqeqfqgqhqiqjqkqlqmqnqoqpqqqrqsqtquqvqwqxqyqzqAqBqCqDqEqFqGqHqIqJqKqLqMqNqOqPqQqRqSqTqUqVqWqXqYqZq0q1q2q3q4q5q6q7q8q9q!q#q$q%q'q(q)q*q+q,q-q.q/q:q;q=q?q@q[q]q^qUc_q`q{q|q}q~qarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzrArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr0r1r2r3r4r5r6r7r8r9r!r#r$r%r'r(r)r*r+rr ,r-r.rn /rs Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pc-b:r.b;r/b=r:b?r;b@r=b[r?b]r@b^r[b_r]b`r^b{r_b|r`b}r{b~r|bas}bbsy ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnccsdsesfsgshsisjsksY lsocmsnsospspcqsvbrssstsusjbvskbwslbxsmbysI zsnbAsobBspbCsDsEsFsGsHsIsJsKsLsMsNsOsPsQsRsSsTsqcUsVsWsrcXsYsZs0s1s2s3s4s5s6s7s

141 context = debug_key_reuse(False) 6 ctx1deK2nx

142 else: 

143 context = nullcontext() 934 ctx2J - d Gbe qbZ . K Hb2 IbTcNkOkk a o / JbKb: LbMbL ; b Nbl ObPbPk0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw QkT RkU SkTkUk} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW VkX Wkcb+b,bv Xkrbsbtb6 ubYkZk0k1k2k3k4k5k6k7k8k9k!k#k$k%k'k(k)k*k+k,k-k.k/k:k;k=k?k@k[k]k^k_k`k{k|k}k~kalblcldlelflglhliljlklllmlnlolplqlrlsltlulvlwlxlylzlAlBlClDlElFlGlHlIlJlKlLlMlNlOlPlQlRlSlTlUlVlWlXlYlZl0l1l2l3l4l5l6l7l8l9l!l#l$l%l'l(l)l*l+l,l-l.l/l:l;l=l?l@l[l]l^l_l`l{l|l}l~lambmcmdmemfmgmhmimjmkmlmmmnmompmqmrmsmtmumvmwmxmymzmAmBmCmDmEmFmGmHmImJmKmLmMmNmOmPmQmRmSmTmUmVmWmXmYmZm0m1m2m3m4m5m6m7m8m9m!m#m$m%m'm(m)m*m+m,m-m.m/m:m;m=m?m@m[m]m^m_m`m{m|m}m~manbncndnenfngnhninjnknlnmnnnonpnqnrnsntnunvnwnxnynznAnBnCnDnEnFnGnHnInJnKnLnMnNnOnPnQnRnSnTnUnVnWnXnYnZn0n1n2n3n4n5n6n7n8n9n!n#n$n%n'n(n)n*n+n,n-n.n/n:n;n=n?n@n[n]n^n_n`n{n|n}n~naobocodoeofogohoiojokolomonooopoqorosotouovowoxoyozoAoBoCoDoEoFoGoHoIoJoKoLoMoNoOoPoQoRoSoToUoVoWoXoYoZo0o1o2o3o4o5o6o7o8o9o!o#o$o%o'o(o)o*o+o,o-o.o/o:o;o=o?o@o[o]o^o_o`o{o|o}o~oapbpcpdpepfpgphpipjpkplpmpnpopppqprpsptpupvpwpxpypzpApBpCpDpEpFpGpHpIpJpKpLpMpNpOpPpQpRpSpTpUpVpWpXpYpZp0p1p2p3p4p5p6p7p8p9p!p#p$p%p'p(p)p*p+p,p-p.p/p:p;p=p?p@p[p]p^p_p`p{p|p}p~paqbqcqdqeqfqgqhqiqjqkqlqmqnqoqpqqqrqsqtquqvqwqxqyqzqAqBqCqDqEqFqGqHqIqJqKqLqMqNqOqPqQqRqSqTqUqVqWqXqYqZq0q1q2q3q4q5q6q7q8q9q!q#q$q%q'q(q)q*q+q,q-q.q/q:q;q=q?q@q[q]q^qUc_q`q{q|q}q~qarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzrArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr0r1r2r3r4r5r6r7r8r9r!r#r$r%r'r(r)r*r+rr ,r-r.rn /rs Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pc-b:r.b;r/b=r:b?r;b@r=b[r?b]r@b^r[b_r]b`r^b{r_b|r`b}r{b~r|bas}bbsy ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnccsdsesfsgshsisjsksY lsocmsnsospspcqsvbrssstsusjbvskbwslbxsmbysI zsnbAsobBspbCsDsEsFsGsHsIsJsKsLsMsNsOsPsQsRsSsTsqcUsVsWsrcXsYsZs0s1s2s3s4s5s6s7s

144 with context: 934 ctx2J - d Gbe qbZ . K Hb2 IbTcNkOkk a o / JbKb: LbMbL ; b Nbl ObPbPk0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw QkT RkU SkTkUk} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW VkX Wkcb+b,bv Xkrbsbtb6 ubYkZk0k1k2k3k4k5k6k7k8k9k!k#k$k%k'k(k)k*k+k,k-k.k/k:k;k=k?k@k[k]k^k_k`k{k|k}k~kalblcldlelflglhliljlklllmlnlolplqlrlsltlulvlwlxlylzlAlBlClDlElFlGlHlIlJlKlLlMlNlOlPlQlRlSlTlUlVlWlXlYlZl0l1l2l3l4l5l6l7l8l9l!l#l$l%l'l(l)l*l+l,l-l.l/l:l;l=l?l@l[l]l^l_l`l{l|l}l~lambmcmdmemfmgmhmimjmkmlmmmnmompmqmrmsmtmumvmwmxmymzmAmBmCmDmEmFmGmHmImJmKmLmMmNmOmPmQmRmSmTmUmVmWmXmYmZm0m1m2m3m4m5m6m7m8m9m!m#m$m%m'm(m)m*m+m,m-m.m/m:m;m=m?m@m[m]m^m_m`m{m|m}m~manbncndnenfngnhninjnknlnmnnnonpnqnrnsntnunvnwnxnynznAnBnCnDnEnFnGnHnInJnKnLnMnNnOnPnQnRnSnTnUnVnWnXnYnZn0n1n2n3n4n5n6n7n8n9n!n#n$n%n'n(n)n*n+n,n-n.n/n:n;n=n?n@n[n]n^n_n`n{n|n}n~naobocodoeofogohoiojokolomonooopoqorosotouovowoxoyozoAoBoCoDoEoFoGoHoIoJoKoLoMoNoOoPoQoRoSoToUoVoWoXoYoZo0o1o2o3o4o5o6o7o8o9o!o#o$o%o'o(o)o*o+o,o-o.o/o:o;o=o?o@o[o]o^o_o`o{o|o}o~oapbpcpdpepfpgphpipjpkplpmpnpopppqprpsptpupvpwpxpypzpApBpCpDpEpFpGpHpIpJpKpLpMpNpOpPpQpRpSpTpUpVpWpXpYpZp0p1p2p3p4p5p6p7p8p9p!p#p$p%p'p(p)p*p+p,p-p.p/p:p;p=p?p@p[p]p^p_p`p{p|p}p~paqbqcqdqeqfqgqhqiqjqkqlqmqnqoqpqqqrqsqtquqvqwqxqyqzqAqBqCqDqEqFqGqHqIqJqKqLqMqNqOqPqQqRqSqTqUqVqWqXqYqZq0q1q2q3q4q5q6q7q8q9q!q#q$q%q'q(q)q*q+q,q-q.q/q:q;q=q?q@q[q]q^qUc_q`q{q|q}q~qarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzrArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr0r1r2r3r4r5r6r7r8r9r!r#r$r%r'r(r)r*r+rr ,r-r.rn /rs Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pc-b:r.b;r/b=r:b?r;b@r=b[r?b]r@b^r[b_r]b`r^b{r_b|r`b}r{b~r|bas}bbsy ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnccsdsesfsgshsisjsksY lsocmsnsospspcqsvbrssstsusjbvskbwslbxsmbysI zsnbAsobBspbCsDsEsFsGsHsIsJsKsLsMsNsOsPsQsRsSsTsqcUsVsWsrcXsYsZs0s1s2s3s4s5s6s7s

145 # jitted-vmapped key split seems to be triggering a false positive 

146 # with key reuse checks 

147 self._keys = _split_unpack(key, num) 934 ctx2J - d Gbe qbZ . K Hb2 IbTcNkOkk a o / JbKb: LbMbL ; b Nbl ObPbPk0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw QkT RkU SkTkUk} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW VkX Wkcb+b,bv Xkrbsbtb6 ubYkZk0k1k2k3k4k5k6k7k8k9k!k#k$k%k'k(k)k*k+k,k-k.k/k:k;k=k?k@k[k]k^k_k`k{k|k}k~kalblcldlelflglhliljlklllmlnlolplqlrlsltlulvlwlxlylzlAlBlClDlElFlGlHlIlJlKlLlMlNlOlPlQlRlSlTlUlVlWlXlYlZl0l1l2l3l4l5l6l7l8l9l!l#l$l%l'l(l)l*l+l,l-l.l/l:l;l=l?l@l[l]l^l_l`l{l|l}l~lambmcmdmemfmgmhmimjmkmlmmmnmompmqmrmsmtmumvmwmxmymzmAmBmCmDmEmFmGmHmImJmKmLmMmNmOmPmQmRmSmTmUmVmWmXmYmZm0m1m2m3m4m5m6m7m8m9m!m#m$m%m'm(m)m*m+m,m-m.m/m:m;m=m?m@m[m]m^m_m`m{m|m}m~manbncndnenfngnhninjnknlnmnnnonpnqnrnsntnunvnwnxnynznAnBnCnDnEnFnGnHnInJnKnLnMnNnOnPnQnRnSnTnUnVnWnXnYnZn0n1n2n3n4n5n6n7n8n9n!n#n$n%n'n(n)n*n+n,n-n.n/n:n;n=n?n@n[n]n^n_n`n{n|n}n~naobocodoeofogohoiojokolomonooopoqorosotouovowoxoyozoAoBoCoDoEoFoGoHoIoJoKoLoMoNoOoPoQoRoSoToUoVoWoXoYoZo0o1o2o3o4o5o6o7o8o9o!o#o$o%o'o(o)o*o+o,o-o.o/o:o;o=o?o@o[o]o^o_o`o{o|o}o~oapbpcpdpepfpgphpipjpkplpmpnpopppqprpsptpupvpwpxpypzpApBpCpDpEpFpGpHpIpJpKpLpMpNpOpPpQpRpSpTpUpVpWpXpYpZp0p1p2p3p4p5p6p7p8p9p!p#p$p%p'p(p)p*p+p,p-p.p/p:p;p=p?p@p[p]p^p_p`p{p|p}p~paqbqcqdqeqfqgqhqiqjqkqlqmqnqoqpqqqrqsqtquqvqwqxqyqzqAqBqCqDqEqFqGqHqIqJqKqLqMqNqOqPqQqRqSqTqUqVqWqXqYqZq0q1q2q3q4q5q6q7q8q9q!q#q$q%q'q(q)q*q+q,q-q.q/q:q;q=q?q@q[q]q^qUc_q`q{q|q}q~qarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzrArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr0r1r2r3r4r5r6r7r8r9r!r#r$r%r'r(r)r*r+rr ,r-r.rn /rs Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pc-b:r.b;r/b=r:b?r;b@r=b[r?b]r@b^r[b_r]b`r^b{r_b|r`b}r{b~r|bas}bbsy ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnccsdsesfsgshsisjsksY lsocmsnsospspcqsvbrssstsusjbvskbwslbxsmbysI zsnbAsobBspbCsDsEsFsGsHsIsJsKsLsMsNsOsPsQsRsSsTsqcUsVsWsrcXsYsZs0s1s2s3s4s5s6s7s

148 self._num_used = 0 934 ctx2J - d Gbe qbZ . K Hb2 IbTcNkOkk a o / JbKb: LbMbL ; b Nbl ObPbPk0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw QkT RkU SkTkUk} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW VkX Wkcb+b,bv Xkrbsbtb6 ubYkZk0k1k2k3k4k5k6k7k8k9k!k#k$k%k'k(k)k*k+k,k-k.k/k:k;k=k?k@k[k]k^k_k`k{k|k}k~kalblcldlelflglhliljlklllmlnlolplqlrlsltlulvlwlxlylzlAlBlClDlElFlGlHlIlJlKlLlMlNlOlPlQlRlSlTlUlVlWlXlYlZl0l1l2l3l4l5l6l7l8l9l!l#l$l%l'l(l)l*l+l,l-l.l/l:l;l=l?l@l[l]l^l_l`l{l|l}l~lambmcmdmemfmgmhmimjmkmlmmmnmompmqmrmsmtmumvmwmxmymzmAmBmCmDmEmFmGmHmImJmKmLmMmNmOmPmQmRmSmTmUmVmWmXmYmZm0m1m2m3m4m5m6m7m8m9m!m#m$m%m'm(m)m*m+m,m-m.m/m:m;m=m?m@m[m]m^m_m`m{m|m}m~manbncndnenfngnhninjnknlnmnnnonpnqnrnsntnunvnwnxnynznAnBnCnDnEnFnGnHnInJnKnLnMnNnOnPnQnRnSnTnUnVnWnXnYnZn0n1n2n3n4n5n6n7n8n9n!n#n$n%n'n(n)n*n+n,n-n.n/n:n;n=n?n@n[n]n^n_n`n{n|n}n~naobocodoeofogohoiojokolomonooopoqorosotouovowoxoyozoAoBoCoDoEoFoGoHoIoJoKoLoMoNoOoPoQoRoSoToUoVoWoXoYoZo0o1o2o3o4o5o6o7o8o9o!o#o$o%o'o(o)o*o+o,o-o.o/o:o;o=o?o@o[o]o^o_o`o{o|o}o~oapbpcpdpepfpgphpipjpkplpmpnpopppqprpsptpupvpwpxpypzpApBpCpDpEpFpGpHpIpJpKpLpMpNpOpPpQpRpSpTpUpVpWpXpYpZp0p1p2p3p4p5p6p7p8p9p!p#p$p%p'p(p)p*p+p,p-p.p/p:p;p=p?p@p[p]p^p_p`p{p|p}p~paqbqcqdqeqfqgqhqiqjqkqlqmqnqoqpqqqrqsqtquqvqwqxqyqzqAqBqCqDqEqFqGqHqIqJqKqLqMqNqOqPqQqRqSqTqUqVqWqXqYqZq0q1q2q3q4q5q6q7q8q9q!q#q$q%q'q(q)q*q+q,q-q.q/q:q;q=q?q@q[q]q^qUc_q`q{q|q}q~qarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzrArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr0r1r2r3r4r5r6r7r8r9r!r#r$r%r'r(r)r*r+rr ,r-r.rn /rs Ac7 Bc8 Cc9 Dc! Ec# Fc$ Gc% Hc' Ic( Jc) Kc* LcebMcfbNcgbOcG Pc-b:r.b;r/b=r:b?r;b@r=b[r?b]r@b^r[b_r]b`r^b{r_b|r`b}r{b~r|bas}bbsy ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnccsdsesfsgshsisjsksY lsocmsnsospspcqsvbrssstsusjbvskbwslbxsmbysI zsnbAsobBspbCsDsEsFsGsHsIsJsKsLsMsNsOsPsQsRsSsTsqcUsVsWsrcXsYsZs0s1s2s3s4s5s6s7s

149 

150 def __len__(self): (empty)

151 return len(self._keys) - self._num_used 882 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jstttpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

152 

153 def pop(self, shape: int | tuple[int, ...] = ()) -> Key[Array, '*batch {shape}']: (empty)

154 """ 

155 Pop one or more keys from the list. 

156 

157 Parameters 

158 ---------- 

159 shape 

160 The shape of the keys to pop. If empty (default), a single key is 

161 popped and returned. If not empty, the popped key is split and 

162 reshaped to the target shape. 

163 

164 Returns 

165 ------- 

166 The popped keys as a jax array with the requested shape. 

167 

168 Raises 

169 ------ 

170 IndexError 

171 If the list is empty. 

172 """ 

173 if len(self) == 0: 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

174 msg = 'No keys left to pop' 1 ctx1n

175 raise IndexError(msg) 1 ctx1n

176 if not isinstance(shape, tuple): 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

177 shape = (shape,) 9 ctx2n DbEbFbY vbQcRcSc

178 key = self._keys[self._num_used] 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

179 self._num_used += 1 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

180 if shape: 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

181 key = _split_shaped(key, shape) 82 ctx2J d e k a o L b l M f N O 3 P 4 Q h R S i 5 w T U p j q V c m W X v Cbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG y + , z A hbB C D E F H ibx DbEbFbY vbQcRcScjbkblbmbI nbobpb

182 return key 880 ctx2J - d Gbe qbZ . K Hb2 Ibsctcuck a o / JbKb: LbMbL ; b Nbl ObPb0 = 1 QbdbRbM ? f SbN TbO @ Ub3 VbP [ Wb4 XbQ ] h YbR ZbS ^ i 0b5 1b_ 2b3b` 4b5b{ 6b7b| 8b9bw T U xcXc} !b#bp ~ j $bq %bab'b(bV bbc )bm *bW X cb+b,bv YcrbZcsbtb6 ub0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzdAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd0d1d2d3d4d5d6d7d8d9d!d#d$d%d'd(d)d*d+d,d-d.d/d:d;d=d?d@d[d]d^d_d`d{d|d}d~daebecedeeefegeheiejekelemeneoepeqereseteuevewexeyezeAeBeCeDeEeFeGeHeIeJeKeLeMeNeOePeQeReSeTeUeVeWeXeYeZe0e1e2e3e4e5e6e7e8e9e!e#e$e%e'e(e)e*e+e,e-e.e/e:e;e=e?e@e[e]e^e_e`e{e|e}e~eafbfcfdfefffgfhfifjfkflfmfnfofpfqfrfsftfufvfwfxfyfzfAfBfCfDfEfFfGfHfIfJfKfLfMfNfOfPfQfRfSfTfUfVfWfXfYfZf0f1f2f3f4f5f6f7f8f9f!f#f$f%f'f(f)f*f+f,f-f.f/f:f;f=f?f@f[f]f^f_f`f{f|f}f~fagbgcgdgegfggghgigjgkglgmgngogpgqgrgsgtgugvgwgxgygzgAgBgCgDgEgFgGgHgIgJgKgLgMgNgOgPgQgRgSgTgUgVgWgXgYgZg0g1g2g3g4g5g6g7g8g9g!g#g$g%g'g(g)g*g+g,g-g.g/g:g;g=g?g@g[g]g^g_g`g{g|g}g~gahbhchdhehfhghhhihjhkhlhmhnhohphqhrhshthuhvhwhxhyhzhAhBhChDhEhFhGhHhIhJhKhLhMhNhOhPhQhRhShThUhVhWhXhYhZh0h1h2h3h4h5h6h7h8h9h!h#h$h%h'h(h)h*h+h,h-h.h/h:h;h=h?h@h[h]h^h_h`h{h|h}h~haibicidieifigihiiijikiliminioipiqirisitiuiviwixiyiziAiBiCiDiEiFiGiHiIiJiKiLiMiNiOiPiQiRiSiTiUiViWiXiYiZi0i1i2i3i4i5i6i7i8i9i!i#i$i%i'i(i)i*i+i,i-i.i/i:i;i=i?i@i[i]i^i_i`i{i|i}i~iajbjcjdjejfjgjhjijjjkjljmjnjojpjqjrjsjtjujvjwjxjyjzjAjBjCjDjEjFjGjHjIjJjKjLjMjNjOjPjQjRjSjTjUjVjWjXjYjZj0j1j2j3j4j5j6j7j8j9j!j#j$j%j'j(j)j*j+j,jt u r -jCbn s 7 8 9 ! # $ % ' ( ) * ebfbgbG -b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}by ~b+ ac, bcz ccA dchbecB fcC gcD hcE icF jcH kciblcx mcnc.jDb/j:j;jEbFb=j?jY oc@j[jpcvbQcRcScjbkblbmbI nbobpb]j^j_j`j{j|j}j~jakbkckdkekfkgkhkikqcjkkkrclkmknkokpkqkrksktkuk

183 

184 

185@partial(jit, static_argnums=(1,)) (empty)

186def _split_unpack( (empty)

187 key: Key[Array, '*batch'], num: int 

188) -> tuple[Key[Array, '*batch'], ...]: 

189 if key.ndim == 0: 13 ctx2d e qbTck o w p q Ucn s x

190 keys = random.split(key, num) 11 ctx2e qbTck o w p q Ucn s

191 elif key.ndim == 1: 191 ↛ 193line 191 didn't jump to line 193 because the condition on line 191 was always true4 ctx1denx

192 keys = vmap(random.split, in_axes=(0, None), out_axes=1)(key, num) 4 ctx1denx

193 return tuple(keys) 13 ctx2d e qbTck o w p q Ucn s x

194 

195 

196@partial(jit, static_argnums=(1,)) (empty)

197def _split_shaped( (empty)

198 key: Key[Array, '*batch'], shape: tuple[int, ...] 

199) -> Key[Array, '*batch {shape}']: 

200 num = math.prod(shape) 17 ctx2e k a o p q v Cbn s G y DbEbFbY I

201 if key.ndim == 0: 201 ↛ 203line 201 didn't jump to line 203 because the condition on line 201 was always true17 ctx2e k a o p q v Cbn s G y DbEbFbY I

202 keys = random.split(key, num) 17 ctx2e k a o p q v Cbn s G y DbEbFbY I

203 elif key.ndim == 1: 

204 keys = vmap(random.split, in_axes=(0, None))(key, num) 

205 return keys.reshape(*key.shape, *shape) 17 ctx2e k a o p q v Cbn s G y DbEbFbY I

206 

207 

208def truncated_normal_onesided( (empty)

209 key: Key[Array, ''], 

210 shape: Sequence[int], 

211 upper: Bool[Array, '*'], 

212 bound: Float32[Array, '*'], 

213 *, 

214 clip: bool = True, 

215) -> Float32[Array, '*']: 

216 """ 

217 Sample from a one-sided truncated standard normal distribution. 

218 

219 Parameters 

220 ---------- 

221 key 

222 JAX random key. 

223 shape 

224 Shape of output array, broadcasted with other inputs. 

225 upper 

226 True for (-∞, bound], False for [bound, ∞). 

227 bound 

228 The truncation boundary. 

229 clip 

230 Whether to clip the truncated uniform samples to (0, 1) before 

231 transforming them to truncated normal. Intended for debugging purposes. 

232 

233 Returns 

234 ------- 

235 Array of samples from the truncated normal distribution. 

236 """ 

237 # Pseudocode: 

238 # | if upper: 

239 # | if bound < 0: 

240 # | ndtri(uniform(0, ndtr(bound))) = 

241 # | ndtri(ndtr(bound) * u) 

242 # | if bound > 0: 

243 # | -ndtri(uniform(ndtr(-bound), 1)) = 

244 # | -ndtri(ndtr(-bound) + ndtr(bound) * (1 - u)) 

245 # | if not upper: 

246 # | if bound < 0: 

247 # | ndtri(uniform(ndtr(bound), 1)) = 

248 # | ndtri(ndtr(bound) + ndtr(-bound) * (1 - u)) 

249 # | if bound > 0: 

250 # | -ndtri(uniform(0, ndtr(-bound))) = 

251 # | -ndtri(ndtr(-bound) * u) 

252 shape = jnp.broadcast_shapes(shape, upper.shape, bound.shape) 11 ctx1dabfhijctur

253 bound_pos = bound > 0 11 ctx1dabfhijctur

254 ndtr_bound = ndtr(bound) 11 ctx1dabfhijctur

255 ndtr_neg_bound = ndtr(-bound) 11 ctx1dabfhijctur

256 scale = jnp.where(upper, ndtr_bound, ndtr_neg_bound) 11 ctx1dabfhijctur

257 shift = jnp.where(upper, ndtr_neg_bound, ndtr_bound) 11 ctx1dabfhijctur

258 u = random.uniform(key, shape) 11 ctx1dabfhijctur

259 left_u = scale * (1 - u) # ~ uniform in (0, ndtr(±bound)] 11 ctx1dabfhijctur

260 right_u = shift + scale * u # ~ uniform in [ndtr(∓bound), 1) 11 ctx1dabfhijctur

261 truncated_u = jnp.where(upper ^ bound_pos, left_u, right_u) 11 ctx1dabfhijctur

262 if clip: 11 ctx1dabfhijctur

263 # on gpu the accuracy is lower and sometimes u can reach the boundaries 

264 zero = jnp.zeros((), truncated_u.dtype) 10 ctx1dabfhijctu

265 one = jnp.ones((), truncated_u.dtype) 10 ctx1dabfhijctu

266 truncated_u = jnp.clip( 10 ctx1dabfhijctu

267 truncated_u, jnp.nextafter(zero, one), jnp.nextafter(one, zero) 

268 ) 

269 truncated_norm = ndtri(truncated_u) 11 ctx1dabfhijctur

270 return jnp.where(bound_pos, -truncated_norm, truncated_norm) 11 ctx1dabfhijctur

271 

272 

273def get_default_device() -> Device: (empty)

274 """Get the current default JAX device.""" 

275 with ensure_compile_time_eval(): 54 ctx2g - . k / : ; = ? f N @ [ ] ^ _ ` { VcGk| } ~ ycabzcWcbbcbAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcz A B C D E F H rt

276 return jnp.zeros(()).device 54 ctx2g - . k / : ; = ? f N @ [ ] ^ _ ` { VcGk| } ~ ycabzcWcbbcbAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcz A B C D E F H rt

277 

278 

279def get_device_count() -> int: (empty)

280 """Get the number of available devices on the default platform.""" 

281 device = get_default_device() 32 ctx2- . k / : ; = ? @ [ ] ^ _ ` { | } ~ ycabzcWcbbcbz A B C D E F H

282 return device_count(device.platform) 32 ctx2- . k / : ; = ? @ [ ] ^ _ ` { | } ~ ycabzcWcbbcbz A B C D E F H

283 

284 

285def is_key(x: object) -> bool: (empty)

286 """Determine if `x` is a jax random key.""" 

287 return isinstance(x, Array) and jnp.issubdtype(x.dtype, prng_key) 94 ctx2J d e Z K 2 k a o vkwk?sxkyk@sL b l 0 1 dbM f N O zk3 P Ak4 Q h R S i 5 BkCk[sDkEk]sFkVcGkHkIk^sw T U xcJk_sp j q yczcWcV c m W X KkLk`sv MkCbs 7 8 9 ! # $ % ' ( ) * y + , z A B C D E F x