Coverage for src/bartz/_jaxext/_jit.py: 92%

1# bartz/src/bartz/_jaxext/_jit.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"""Signature-preserving `jax.jit` wrapper.""" 

26 

27from collections.abc import Callable, Sequence (empty)

28from typing import ( (empty)

29 TYPE_CHECKING, 

30 Any, 

31 Concatenate, 

32 ParamSpec, 

33 Protocol, 

34 TypeVar, 

35 overload, 

36 runtime_checkable, 

37) 

38 

39from jax import ShapeDtypeStruct (empty)

40from jax import jit as _jax_jit (empty)

41from jax.stages import Lowered, Traced (empty)

42from jaxtyping import PyTree (empty)

43 

44_P = ParamSpec('_P') (empty)

45_P2 = ParamSpec('_P2') (empty)

46_R = TypeVar('_R') (empty)

47_R2 = TypeVar('_R2') (empty)

48_R_co = TypeVar('_R_co', covariant=True) (empty)

49_T = TypeVar('_T') (empty)

50 

51 

52@runtime_checkable (empty)

53class JitWrapped(Protocol[_P, _R_co]): (empty)

54 """Static type of a jitted function: the wrapped signature plus jit methods.""" 

55 

56 def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R_co: ... (empty)

57 

58 def clear_cache(self) -> None: ... (empty)

59 

60 def eval_shape( (empty)

61 self, *args: _P.args, **kwargs: _P.kwargs (empty)

62 ) -> PyTree[ShapeDtypeStruct]: ... (empty)

63 

64 def lower(self, *args: _P.args, **kwargs: _P.kwargs) -> Lowered: ... (empty)

65 

66 def trace(self, *args: _P.args, **kwargs: _P.kwargs) -> Traced: ... (empty)

67 

68 # jax's jitted callables implement the descriptor protocol, so `@jit` also 

69 # works on methods; declare it so type checkers bind `self` on attribute 

70 # access (the bound object proxies the jit methods through `__func__`, but 

71 # their bound signatures are not remapped -- only `__call__` is) 

72 

73 @overload (empty)

74 def __get__( (empty)

75 self, obj: None, objtype: type | None = None, / (empty)

76 ) -> 'JitWrapped[_P, _R_co]': ... (empty)

77 

78 @overload (empty)

79 def __get__( (empty)

80 self: 'JitWrapped[Concatenate[_T, _P2], _R2]', (empty)

81 obj: _T, (empty)

82 objtype: type | None = None, (empty)

83 /, 

84 ) -> Callable[_P2, _R2]: ... (empty)

85 

86 if not TYPE_CHECKING: 86 ↛ anywhereline 86 didn't jump anywhere: it always raised an exception.(empty)

87 # WORKAROUND(beartype<99): beartype chokes on ParamSpec-subscripted 

88 # generics, and the jaxtyping import hook used by the test suite makes 

89 # it process the `JitWrapped[_P, _R]` hints in `jit`'s overloads. Erase 

90 # the subscript at runtime so beartype sees the plain runtime-checkable 

91 # protocol, which jitted functions genuinely satisfy. `99` is a 

92 # placeholder for the beartype release gaining PEP 612 generics support. 

93 def __class_getitem__(cls, item: object) -> type: (empty)

94 return cls (empty)

95 

96 

97# WORKAROUND(jax<99): `jax.jit` is typed to return `JitWrapped`, which erases the 

98# wrapped function's signature, so static checkers can't validate calls to jitted 

99# functions. This shim recovers the signature via `ParamSpec`, declaring our own 

100# `JitWrapped` protocol that combines it with the jit-specific methods (including 

101# `clear_cache`, which jax adds to the jitted callable at runtime and omits from 

102# its own static `JitWrapped` type). Tracked upstream at jax-ml/jax#23719; the 

103# jax maintainers are blocked on migrating internal Google code to a type checker 

104# that understands `ParamSpec` (jax itself has moved to pyrefly). Once `jax.jit` 

105# preserves the signature natively, this whole module can go and `jit` can be 

106# imported straight from jax. `99` is a placeholder for that unknown future jax 

107# release. 

108@overload (empty)

109def jit( (empty)

110 fun: Callable[_P, _R], (empty)

111 /, 

112 *, 

113 static_argnums: int | Sequence[int] | None = ..., (empty)

114 static_argnames: str | Sequence[str] | None = ..., (empty)

115 donate_argnums: int | Sequence[int] | None = ..., (empty)

116 **kwargs: Any, (empty)

117) -> JitWrapped[_P, _R]: ... (empty)

118 

119 

120@overload (empty)

121def jit( (empty)

122 fun: None = ..., (empty)

123 /, 

124 *, 

125 static_argnums: int | Sequence[int] | None = ..., (empty)

126 static_argnames: str | Sequence[str] | None = ..., (empty)

127 donate_argnums: int | Sequence[int] | None = ..., (empty)

128 **kwargs: Any, (empty)

129) -> Callable[[Callable[_P, _R]], JitWrapped[_P, _R]]: ... (empty)

130def jit(fun: Any = None, /, **kwargs: Any) -> Any: (empty)

131 """Wrap `jax.jit` preserving the wrapped function's static type signature. 

132 

133 `jax.jit` is typed to return an opaque ``JitWrapped`` callable, which erases 

134 the wrapped signature; static checkers then treat every call to a jitted 

135 function as returning an unknown type, cascading into false positives. This 

136 shim is typed with a `ParamSpec` so jitted calls keep their real signature 

137 and argument checking, while at runtime it just defers to `jax.jit`. 

138 

139 Use it as a drop-in for both decorator forms, ``@jit`` and ``@jit(...)``. 

140 

141 Parameters 

142 ---------- 

143 fun 

144 The function to compile, or `None` to use the keyword-only form. 

145 **kwargs 

146 Keyword arguments forwarded to `jax.jit` (e.g. `static_argnums`, 

147 `static_argnames`, `donate_argnums`). 

148 

149 Returns 

150 ------- 

151 The jitted function, or a decorator if `fun` is `None`. 

152 """ 

153 # WORKAROUND(jax<0.8.1): jax gained native `@jit(...)` two-stage decorator 

154 # support in 0.8.1. Once the floor reaches 0.8.1 the runtime fallback could 

155 # defer to jax's native form, but keep the shim regardless, because jax's 

156 # own overloads still return `JitWrapped` and erase the signature; the 

157 # ParamSpec typing here is the whole point. 

158 if fun is None: (empty)

159 return lambda f: _jax_jit(f, **kwargs) (empty)

160 else: 

161 return _jax_jit(fun, **kwargs) (empty)