Coverage for src/bartz/mcmcstep/_lazy.py: 98%

1# bartz/src/bartz/mcmcstep/_lazy.py 

2# 

3# Copyright (c) 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"""Deferred array construction used to lay out the MCMC state before sharding.""" 

26 

27from collections.abc import Callable (empty)

28from typing import Any, TypeVar, cast, overload (empty)

29 

30from equinox import Module (empty)

31from jax import ShapeDtypeStruct, tree (empty)

32from jax import numpy as jnp (empty)

33from jax.typing import DTypeLike (empty)

34from jaxtyping import Array, PyTree, Shaped (empty)

35 

36T = TypeVar('T') (empty)

37 

38 

39class _LazyArray(Module): (empty)

40 """Like `functools.partial` but specialized to array-creating functions like `jax.numpy.zeros`.""" 

41 

42 array_creator: Callable 

43 shape: tuple[int, ...] 

44 args: tuple 

45 

46 def __init__( (empty)

47 self, array_creator: Callable, shape: tuple[int, ...], *args: Any 

48 ) -> None: 

49 self.array_creator = array_creator (empty)

50 self.shape = shape (empty)

51 self.args = args (empty)

52 

53 def __call__(self, **kwargs: Any) -> T: (empty)

54 return self.array_creator(self.shape, *self.args, **kwargs) (empty)

55 

56 @property (empty)

57 def ndim(self) -> int: (empty)

58 return len(self.shape) (empty)

59 

60 @property (empty)

61 def dtype(self) -> DTypeLike: (empty)

62 # The concrete dtype is unknown until the array is built; report the 

63 # abstract `generic` scalar type so jaxtyping's `Shaped[_LazyArray, ...]` 

64 # runtime check can read `.dtype` (it ignores the dtype name anyway). 

65 return jnp.generic 

66 

67 

68DummyArray = Array | ShapeDtypeStruct | _LazyArray (empty)

69 

70 

71# WORKAROUND(jaxtyping<0.3.9): a shared structure variable 

72# (PyTree[DummyArray, 'T'] -> PyTree[ShapeDtypeStruct, 'T']) is mis-bound to a 

73# single leaf when the leaf type is a union containing a Module (here 

74# `_LazyArray`), so the return-value check spuriously fails. Drop the structure 

75# variable; `tree.map` preserves the structure anyway. Restore 'T' on both 

76# annotations once the jaxtyping floor reaches 0.3.9. 

77def add_dummy_axis(x: PyTree[DummyArray]) -> PyTree[ShapeDtypeStruct]: (empty)

78 """Replace array-like leaves with a rank-inflated placeholder.""" 

79 

80 def replace_leaf(leaf: Shaped[DummyArray, '...']) -> ShapeDtypeStruct: (empty)

81 return ShapeDtypeStruct((0,) * (leaf.ndim + 1), jnp.float32) (empty)

82 

83 return tree.map(replace_leaf, x, is_leaf=lambda x: isinstance(x, _LazyArray)) (empty)

84 

85 

86def _lazy( (empty)

87 array_creator: Callable, shape: tuple[int, ...], *args: Any 

88) -> Shaped[Array, '...']: 

89 """Build a `_LazyArray` placeholder, typed as the `Array` it stands in for. 

90 

91 The placeholder is parked in an array-typed state field until `init` 

92 concretizes and shards it, so `cast` hides the deliberate type mismatch 

93 from the static checker (the runtime check is disabled meanwhile). 

94 """ 

95 return cast(Array, _LazyArray(array_creator, shape, *args)) (empty)

96 

97 

98def _return_array( (empty)

99 shape: tuple[int, ...], # noqa: ARG001 

100 arr: Shaped[Array, '*shape'], 

101 **kwargs: Any, # noqa: ARG001 

102) -> Shaped[Array, '*shape']: 

103 """`_LazyArray` factory that returns an already-built array.""" 

104 return arr (empty)

105 

106 

107@overload (empty)

108def _lazy_from_array(arr: Shaped[Array, '*shape']) -> Shaped[Array, '*shape']: ... (empty)

109 

110 

111@overload (empty)

112def _lazy_from_array(arr: None) -> None: ... (empty)

113 

114 

115def _lazy_from_array( (empty)

116 arr: Shaped[Array, '*shape'] | None, 

117) -> Shaped[Array, '*shape'] | None: 

118 """Wrap an existing array as a `_LazyArray` reporting `arr.shape`, or pass `None`.""" 

119 if arr is None: (empty)

120 return None (empty)

121 return _lazy(_return_array, arr.shape, arr) (empty)

122 

123 

124def _broadcast_chain( (empty)

125 shape: tuple[int, ...], 

126 inner: Shaped[_LazyArray, '...'], 

127 chain_axis: int, 

128 **kwargs: Any, 

129) -> Shaped[Array, '...']: 

130 """Concretize `inner` then insert and broadcast a chain axis at `chain_axis`.""" 

131 arr = inner(**kwargs) (empty)

132 arr = jnp.expand_dims(arr, chain_axis) (empty)

133 return jnp.broadcast_to(arr, shape) (empty)

134 

135 

136def _wrap_chain( (empty)

137 inner: Shaped[_LazyArray, '...'], chain_axis: int | None, num_chains: int | None 

138) -> Shaped[_LazyArray, '...']: 

139 """Wrap `inner` so its factory inserts and broadcasts `num_chains` at `chain_axis`. No-op when `chain_axis` is `None`.""" 

140 if chain_axis is None: (empty)

141 return inner (empty)

142 assert num_chains is not None (empty)

143 new_shape = (*inner.shape[:chain_axis], num_chains, *inner.shape[chain_axis:]) (empty)

144 return _LazyArray(_broadcast_chain, new_shape, inner, chain_axis) (empty)

145 

146 

147def _is_lazy_or_none(x: object) -> bool: (empty)

148 """`tree.map(is_leaf=...)` predicate that stops at `_LazyArray` or `None`.""" 

149 return x is None or isinstance(x, _LazyArray) (empty)