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

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

2# 

3# Copyright (c) 2024-2025, 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 functools import partial (empty)

30 

31import jax (empty)

32from jax import jit, random (empty)

33from jax import numpy as jnp (empty)

34from jax.lax import scan (empty)

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

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

37 

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

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

40 

41 

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

43 """ 

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

45 

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

47 arguments have additional axes due to vmap. 

48 """ 

49 doc = fun.__doc__ (empty)

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

51 fun.__doc__ = doc (empty)

52 return fun (empty)

53 

54 

55def minimal_unsigned_dtype(value): (empty)

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

57 if value < 2**8: 97 ctx2w d L M N 2 3 4 5 dcecCcb fco gca Dc6 hcicEcjckcFcx e l O f 7 y g P z lcQ A mcR B h C D i S ncocGcpcqcHcrcT IcsctcJcu E F ' ucKcr j G U vcLcH k m I J wcxcMcs yc8 9 ! V # v W NcX OcY Z Pc0 1 QcRc

58 return jnp.uint8 97 ctx2w d L M N 2 3 4 5 dcecCcb fco gca Dc6 hcicEcjckcFcx e l O f 7 y g P z lcQ A mcR B h C D i S ncocGcpcqcHcrcT IcsctcJcu E F ' ucKcr j G U vcLcH k m I J wcxcMcs yc8 9 ! V # v W NcX OcY Z Pc0 1 QcRc

59 if value < 2**16: 59 ↛ 61line 59 didn't jump to line 61 because the condition on line 59 was always true59 ctx2w d M N 3 4 5 dcecb fco gca hcicjckcx e O f y g z lcA mcB h C D i ncocpcqcrcT sctcu E F ucr j G U vcH k I J wcxcs ycV

60 return jnp.uint16 59 ctx2w d M N 3 4 5 dcecb fco gca hcicjckcx e O f y g z lcA mcB h C D i ncocpcqcrcT sctcu E F ucr j G U vcH k I J wcxcs ycV

61 if value < 2**32: 

62 return jnp.uint32 

63 return jnp.uint64 

64 

65 

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

67def unique( (empty)

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

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

70 """ 

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

72 

73 Parameters 

74 ---------- 

75 x 

76 The input array. 

77 size 

78 The length of the output. 

79 fill_value 

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

81 of unique values in `x`. 

82 

83 Returns 

84 ------- 

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

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

87 actual_length : int 

88 The number of used values in `out`. 

89 """ 

90 if x.size == 0: 14 ctx2b a l m ScBc$ % W X Y Z 0 1

91 return jnp.full(size, fill_value, x.dtype), 0 1 ctx2Sc

92 if size == 0: 13 ctx2b a l m Bc$ % W X Y Z 0 1

93 return jnp.empty(0, x.dtype), 0 1 ctx2Bc

94 x = jnp.sort(x) 12 ctx1balm$%WXYZ01

95 

96 def loop(carry, x): 12 ctx1balm$%WXYZ01

97 i_out, last, out = carry 12 ctx1balm$%WXYZ01

98 i_out = jnp.where(x == last, i_out, i_out + 1) 12 ctx1balm$%WXYZ01

99 out = out.at[i_out].set(x) 12 ctx1balm$%WXYZ01

100 return (i_out, x, out), None 12 ctx1balm$%WXYZ01

101 

102 carry = 0, x[0], jnp.full(size, fill_value, x.dtype) 12 ctx1balm$%WXYZ01

103 (actual_length, _, out), _ = scan(loop, carry, x[:size]) 12 ctx1balm$%WXYZ01

104 return out, actual_length + 1 12 ctx1balm$%WXYZ01

105 

106 

107class split: (empty)

108 """ 

109 Split a key into `num` keys. 

110 

111 Parameters 

112 ---------- 

113 key 

114 The key to split. 

115 num 

116 The number of keys to split into. 

117 """ 

118 

119 _keys: tuple[Key[Array, ''], ...] 

120 _num_used: int 

121 

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

123 self._keys = _split_unpack(key, num) 188 ctx2w ( d ) L * M + N , 2 - zcTcUcb o a 6 . / : ; = ? x @ e [ l ] ^ VcO _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu WcE XcF YcZc0cybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI 1cJ 2cKbLbMbs 3c8 9 ! V # 4cAc5c6c7c8c9c!c#c$c%c'c(cn )c*ct +cv ,cNb-c.c/cK :cOb;c=c?c@cPb[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdQbidRbjdkdldSbmdndodpdqdrdsdtdudvdwd

124 self._num_used = 0 188 ctx2w ( d ) L * M + N , 2 - zcTcUcb o a 6 . / : ; = ? x @ e [ l ] ^ VcO _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu WcE XcF YcZc0cybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI 1cJ 2cKbLbMbs 3c8 9 ! V # 4cAc5c6c7c8c9c!c#c$c%c'c(cn )c*ct +cv ,cNb-c.c/cK :cOb;c=c?c@cPb[c]c^c_c`c{c|c}c~cadbdcdddedfdgdhdQbidRbjdkdldSbmdndodpdqdrdsdtdudvdwd

125 

126 def __len__(self): (empty)

127 return len(self._keys) - self._num_used 174 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bxdydPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

128 

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

130 """ 

131 Pop one or more keys from the list. 

132 

133 Parameters 

134 ---------- 

135 shape 

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

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

138 reshaped to the target shape. 

139 

140 Returns 

141 ------- 

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

143 

144 Raises 

145 ------ 

146 IndexError 

147 If the list is empty. 

148 """ 

149 if len(self) == 0: 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

150 msg = 'No keys left to pop' 1 ctx1t

151 raise IndexError(msg) 1 ctx1t

152 if not isinstance(shape, tuple): 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

153 shape = (shape,) 38 ctx1wdLboaxelfygPzQARBhCDiSuEFrjGHkmIJstvK

154 key = self._keys[self._num_used] 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

155 self._num_used += 1 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

156 if shape: 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

157 key = _split_shaped(key, shape) 38 ctx1wdLboaxelfygPzQARBhCDiSuEFrjGHkmIJstvK

158 return key 172 ctx2w ( d ) L * M + N , 2 - 3 4 5 b o a 6 . / : ; = ? x @ e [ l ] ^ O _ f ` 7 { y | g } P ~ z abbbQ cbA dbebR fbB gbh hbC ibD jbi kbS lbmbnbobpbqbrbsbT tbubvbwbxbu E F ' TbybzbAbr Bbj CbG DbU EbFbGbH Hbk Ibm JbI J KbLbMbs Ub8 Vb9 ! V # WbXbYbZb0b1b2b3b4b5b6bp q n 7bt v Nb8b9bK Ob!b#bPb$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@bQbRb[b]bSb^b_b`b{b|b}b~bacbccc

159 

160 

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

162def _split_unpack(key: Key[Array, ''], num: int) -> tuple[Key[Array, ''], ...]: (empty)

163 keys = random.split(key, num) 6 ctx2zcb o u Act

164 return tuple(keys) 6 ctx2zcb o u Act

165 

166 

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

168def _split_shaped(key: Key[Array, ''], shape: tuple[int, ...]) -> Key[Array, '*']: (empty)

169 num = math.prod(shape) 8 ctx1boarstvK

170 keys = random.split(key, num) 8 ctx1boarstvK

171 return keys.reshape(shape) 8 ctx1boarstvK

172 

173 

174def truncated_normal_onesided( (empty)

175 key: Key[Array, ''], 

176 shape: Sequence[int], 

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

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

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

180 """ 

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

182 

183 Parameters 

184 ---------- 

185 key 

186 JAX random key. 

187 shape 

188 Shape of output array, broadcasted with other inputs. 

189 upper 

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

191 bound 

192 The truncation boundary. 

193 

194 Returns 

195 ------- 

196 Array of samples from the truncated normal distribution. 

197 """ 

198 # Pseudocode: 

199 # | if upper: 

200 # | if bound < 0: 

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

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

203 # | if bound > 0: 

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

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

206 # | if not upper: 

207 # | if bound < 0: 

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

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

210 # | if bound > 0: 

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

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

213 shape = jnp.broadcast_shapes(shape, upper.shape, bound.shape) 12 ctx1daefghijkpqn

214 bound_pos = bound > 0 12 ctx1daefghijkpqn

215 ndtr_bound = ndtr(bound) 12 ctx1daefghijkpqn

216 ndtr_neg_bound = ndtr(-bound) 12 ctx1daefghijkpqn

217 scale = jnp.where(upper, ndtr_bound, ndtr_neg_bound) 12 ctx1daefghijkpqn

218 shift = jnp.where(upper, ndtr_neg_bound, ndtr_bound) 12 ctx1daefghijkpqn

219 u = random.uniform(key, shape) 12 ctx1daefghijkpqn

220 left_u = scale * (1 - u) # ~ uniform in (0, ndtr(±bound)] 12 ctx1daefghijkpqn

221 right_u = shift + scale * u # ~ uniform in [ndtr(∓bound), 1) 12 ctx1daefghijkpqn

222 truncated_u = jnp.where(upper ^ bound_pos, left_u, right_u) 12 ctx1daefghijkpqn

223 truncated_norm = ndtri(truncated_u) 12 ctx1daefghijkpqn

224 return jnp.where(bound_pos, -truncated_norm, truncated_norm) 12 ctx1daefghijkpqn