Coverage for src/lsqfitgp/_patch_jax.py: 88%

1# lsqfitgp/_patch_jax.py 

2# 

3# Copyright (c) 2022, 2023, 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 

20"""Modifications to the global state of jax.""" 

21 

22import gvar 3 ctx1cab

23import numpy 3 ctx1cab

24from jax import config, tree_util 3 ctx1cab

25 

26config.update('jax_enable_x64', True) 3 ctx1cab

27 

28 

29class BufferDictPyTreeDef: 3 ctx1cab

30 @staticmethod 3 ctx1cab

31 def _skeleton(bd): 3 ctx1cab

32 """Return a memoryless BufferDict with the same layout as `bd`.""" 

33 return gvar.BufferDict(bd, buf=numpy.empty(bd.buf.shape, [])) 2 ctx1ab

34 # BufferDict mirrors the data type of the _buf attribute, so we do not 

35 # need to preserve it to maintain consistency. buf is not copied. 

36 

37 def __init__(self, bd): 3 ctx1cab

38 self.skeleton = self._skeleton(bd) 2 ctx1ab

39 self.layout = {k: tuple(bd.slice_shape(k)) for k in bd} 2 ctx1ab

40 # it is not necessary to save the data type because that's in buf 

41 

42 def __eq__(self, other): 3 ctx1cab

43 if not isinstance(other, __class__): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true2 ctx1ab

44 return NotImplemented 

45 return self.layout == other.layout 2 ctx1ab

46 

47 def __hash__(self): 3 ctx1cab

48 return hash(self.layout) 

49 

50 def __repr__(self): 3 ctx1cab

51 return repr(self.layout) 

52 

53 @classmethod 3 ctx1cab

54 def flatten(cls, bd): 3 ctx1cab

55 return (bd.buf,), cls(bd) 2 ctx1ab

56 

57 @classmethod 3 ctx1cab

58 def unflatten(cls, self, children): 3 ctx1cab

59 (buf,) = children 2 ctx1ab

60 new = cls._skeleton(self.skeleton) 2 ctx1ab

61 # copy the skeleton to permit multiple unflattening 

62 new._extension = {} 2 ctx1ab

63 new._buf = buf 2 ctx1ab

64 return new 2 ctx1ab

65 

66 

67# register BufferDict as a pytree 

68tree_util.register_pytree_node( 3 ctx1cab

69 gvar.BufferDict, BufferDictPyTreeDef.flatten, BufferDictPyTreeDef.unflatten 

70)