File tree Expand file tree Collapse file tree 1 file changed +1
-17
lines changed Expand file tree Collapse file tree 1 file changed +1
-17
lines changed Original file line number Diff line number Diff line change 1
- """Patches for `quax`.
2
-
3
- TODO: remove this file when `quax` issues are fixed:
4
-
5
- - https://github.com/issues/created?issue=patrick-kidger%7Cquax%7C57
6
-
7
- """
1
+ """Patches for `quax`."""
8
2
9
3
__all__ : list [str ] = []
10
4
18
12
from jaxtyping import Array , ArrayLike
19
13
20
14
21
- @quax .register (lax .regularized_incomplete_beta_p ) # type: ignore[misc]
22
- def regularized_incomplete_beta_p (
23
- a : ArrayLike ,
24
- b : ArrayLike ,
25
- x : ArrayLike ,
26
- ) -> Array :
27
- """Patched implementation regularized incomplete beta function."""
28
- return lax .regularized_incomplete_beta_p .bind (a , b , x )
29
-
30
-
31
15
@quax .register (lax .scan_p ) # type: ignore[misc]
32
16
def scan_p (* args : ArrayLike , ** kw : Any ) -> Array :
33
17
"""Patched implementation of lax.map."""
You can’t perform that action at this time.
0 commit comments