Kernel refactor
I thought it would be nice to take the BlackJax approach, and be explicit about the kernel as a function of type:
```
MCMCState -> (MCMCState, MCMCInfo)
```
More concretely, I have introduced:
```python
class MCLMCState(NamedTuple):
"""State of the MCLMC algorithm.
"""
x: ArrayTree
u: ArrayTree
l: float
g: ArrayTree
key : Any
class MCLMCInfo(NamedTuple):
transformed_x: ArrayTree
l: ArrayTree
de: float
```
I've only refactored for `sample_normal`, but it could easily be extended everywhere.
This should make adding to BlackJax easier, but I think also clarifies the information flow of the code somewhat.
合并状态:未合并 关闭于 2023-11-05 5 条评论