Coverage for src/lsqfitgp/_jaxext/__init__.py: 90%

1# lsqfitgp/_jaxext/__init__.py 

2# 

3# Copyright (c) 2022, 2023, 2025, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20import functools 3 ctx1bca

21import traceback 3 ctx1bca

22 

23import jax 3 ctx1bca

24from jax import numpy as jnp 3 ctx1bca

25 

26from lsqfitgp._jaxext._batcher import batchufunc # noqa: F401 3 ctx1bca

27from lsqfitgp._jaxext._fasthash import fasthash32, fasthash64 # noqa: F401 3 ctx1bca

28 

29 

30def makejaxufunc(ufunc, *derivs, excluded=None, floatcast=False): 3 ctx1bca

31 """ 

32 Wrap a numpy ufunc to add jax support. 

33 

34 Parameters 

35 ---------- 

36 ufunc : callable 

37 Elementwise function following numpy broadcasting and type promotion 

38 rules. Keyword arguments not supported. 

39 *derivs : sequence of callable 

40 Derivatives of the function w.r.t. each positional argument, with the 

41 same signature as `ufunc`. Pass None to indicate a missing derivative. 

42 There must be as many derivatives as the arguments to `ufunc`. 

43 excluded : sequence of int, optional 

44 The indices of arguments that are not broadcasted. 

45 floatcast : bool, default False 

46 If True, cast all arguments to float before calling the ufunc. 

47 

48 Return 

49 ------ 

50 func : callable 

51 Wrapped `ufunc`. Supports jit, but the calculation is performed on cpu. 

52 """ 

53 nondiff_argnums = [i for i, d in enumerate(derivs) if d is None] 3 ctx1bca

54 

55 @functools.wraps(ufunc) 3 ctx1bca

56 @functools.partial(jax.custom_jvp, nondiff_argnums=nondiff_argnums) 3 ctx1bca

57 def func(*args): 3 ctx1bca

58 args = tuple(map(jnp.asarray, args)) 3 ctx1bca

59 if floatcast: 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true3 ctx1bca

60 flt = float_type(*args) 

61 args = tuple(a.astype(flt) for a in args) 

62 return pure_callback_ufunc( 3 ctx1bca

63 ufunc, jnp.result_type(*args), *args, excluded=excluded 

64 ) 

65 

66 @func.defjvp 3 ctx1bca

67 def func_jvp(*allargs): 3 ctx1bca

68 ndargs = allargs[:-2] 1 ctx1a

69 dargs = allargs[-2] 1 ctx1a

70 dargst = allargs[-1] 1 ctx1a

71 

72 itnd = iter(ndargs) 1 ctx1a

73 itd = iter(dargs) 1 ctx1a

74 args = [next(itnd) if d is None else next(itd) for d in derivs] 1 ctx1a

75 

76 result = func(*args) 1 ctx1a

77 tangent = sum( 1 ctx1a

78 [ 

79 d(*args) * t 

80 for t, d in zip( 

81 dargst, (d for d in derivs if d is not None), strict=True 

82 ) 

83 ] 

84 ) 

85 return result, tangent 1 ctx1a

86 

87 return func 3 ctx1bca

88 

89 

90def elementwise_grad(fun, argnum=0): 3 ctx1bca

91 assert int(argnum) == argnum, argnum 3 ctx1bca

92 assert argnum >= 0, argnum 3 ctx1bca

93 

94 @functools.wraps(fun) 3 ctx1bca

95 def funderiv(*args, **kw): 3 ctx1bca

96 preargs = args[:argnum] 3 ctx1bca

97 postargs = args[argnum + 1 :] 3 ctx1bca

98 

99 def oneargfun(arg): 3 ctx1bca

100 args = (*preargs, arg, *postargs) 3 ctx1bca

101 return fun(*args, **kw) 3 ctx1bca

102 

103 primal = args[argnum] 3 ctx1bca

104 shape = getattr(primal, 'shape', ()) 3 ctx1bca

105 dtype = getattr(primal, 'dtype', type(primal)) 3 ctx1bca

106 tangent = jnp.ones(shape, dtype) 3 ctx1bca

107 _primal_out, tangent_out = jax.jvp(oneargfun, (primal,), (tangent,)) 3 ctx1bca

108 return tangent_out 3 ctx1bca

109 

110 return funderiv 3 ctx1bca

111 

112 

113class skipifabstract: 3 ctx1bca

114 """ 

115 Context manager to do all operations eagerly, or skip them. 

116 

117 It tries to do all operations eagerly even during jit, and skips entirely 

118 if it is not possible. 

119 """ 

120 

121 # I feared this would be slow because of the slow jax exception handling, 

122 # but %timeit suggests it isn't 

123 

124 ENSURE_COMPILE_TIME_EVAL = True 3 ctx1bca

125 ENABLED = True 3 ctx1bca

126 

127 def __enter__(self): 3 ctx1bca

128 if self.ENSURE_COMPILE_TIME_EVAL and self.ENABLED: 128 ↛ exitline 128 didn't return from function '__enter__' because the condition on line 128 was always true3 ctx1bca

129 self.mgr = jax.ensure_compile_time_eval() 3 ctx1bca

130 self.mgr.__enter__() 3 ctx1bca

131 

132 def __exit__(self, exc_type, exc_value, tb): 3 ctx1bca

133 if not self.ENABLED: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true3 ctx1bca

134 return None 

135 suppress = None 3 ctx1bca

136 if self.ENSURE_COMPILE_TIME_EVAL: 136 ↛ 138line 136 didn't jump to line 138 because the condition on line 136 was always true3 ctx1bca

137 suppress = self.mgr.__exit__(exc_type, exc_value, tb) 3 ctx1bca

138 ignorable_error = exc_type is not None and issubclass( 3 ctx1bca

139 exc_type, 

140 (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError), 

141 ) 

142 if suppress or ignorable_error: 3 ctx1bca

143 return True 3 ctx1bca

144 

145 weird_cond = exc_type is IndexError and ( 3 ctx1bca

146 traceback.extract_tb(tb)[-1].name in ('arg_info_pytree', '_origin_msg'), 

147 ) 

148 if weird_cond: # pragma: no cover 3 ctx1bca

149 return True 

150 return None 3 ctx1bca

151 

152 

153def float_type(*args): 3 ctx1bca

154 t = jnp.result_type(*args) 3 ctx1bca

155 return jnp.sin(jnp.empty(0, t)).dtype 3 ctx1bca

156 # numpy does this with common_type, but that supports only arrays, not 

157 # dtypes in the input. jnp.common_type is not defined. 

158 

159 

160def is_jax_type(dtype): 3 ctx1bca

161 dtype = jnp.dtype(dtype) 3 ctx1bca

162 try: 3 ctx1bca

163 jnp.empty(0, dtype) 3 ctx1bca

164 except TypeError as e: 1 ctx1a

165 if 'JAX only supports number' in str(e): 165 ↛ 167line 165 didn't jump to line 167 because the condition on line 165 was always true1 ctx1a

166 return False 1 ctx1a

167 raise 

168 else: 

169 return True 3 ctx1bca

170 

171 

172def pure_callback_ufunc(callback, dtype, *args, excluded=None, **kwargs): 3 ctx1bca

173 """ 

174 Version of jax.pure_callback that deals correctly with ufuncs. 

175 

176 See https://github.com/google/jax/issues/17187. 

177 """ 

178 if excluded is None: 3 ctx1bca

179 excluded = () 3 ctx1bca

180 shape = jnp.broadcast_shapes( 3 ctx1bca

181 *(a.shape for i, a in enumerate(args) if i not in excluded) 

182 ) 

183 ndim = len(shape) 3 ctx1bca

184 padded_args = [ 3 ctx1bca

185 a if i in excluded else jnp.expand_dims(a, tuple(range(ndim - a.ndim))) 

186 for i, a in enumerate(args) 

187 ] 

188 result = jax.ShapeDtypeStruct(shape, dtype) 3 ctx1bca

189 return jax.pure_callback( 3 ctx1bca

190 callback, result, *padded_args, vmap_method='expand_dims', **kwargs 

191 ) 

192 

193 

194def limit_derivatives(x, n, error_func=None): 3 ctx1bca

195 """ 

196 Limit the number of derivatives that goes through a value. 

197 

198 Parameters 

199 ---------- 

200 x : array_like 

201 The value. 

202 n : int 

203 The maximum number of derivatives allowed. Must be an integer. 

204 error_func : callable, optional 

205 A function that takes (derivatives taken, n) and returns an exception. 

206 

207 Return 

208 ------ 

209 x : array_like 

210 The value, unchanged. 

211 """ 

212 assert n == int(n) 3 ctx1bca

213 if error_func is None: 213 ↛ 215line 213 didn't jump to line 215 because the condition on line 213 was never true3 ctx1bca

214 

215 def error_func(current, n): 

216 return ValueError(f'took {current} derivatives > limit {n}') 

217 

218 return _limit_derivatives_impl(0, n, error_func, x) 3 ctx1bca

219 

220 

221@functools.partial(jax.custom_jvp, nondiff_argnums=(0, 1, 2)) 3 ctx1bca

222def _limit_derivatives_impl(current, limit, func, x): 3 ctx1bca

223 if current > limit: 3 ctx1bca

224 raise func(current, limit) 1 ctx1a

225 return x 3 ctx1bca

226 

227 

228@_limit_derivatives_impl.defjvp 3 ctx1bca

229def _limit_derivatives_impl_jvp(current, limit, func, primals, tangents): 3 ctx1bca

230 (x,) = primals 3 ctx1bca

231 (xdot,) = tangents 3 ctx1bca

232 return _limit_derivatives_impl(current + 1, limit, func, x), xdot 3 ctx1bca