Use jaxified logp for initial point evaluation when sampling via Jax#7610

Open

nataziel wants to merge 9 commits intopymc-devs:mainfrom nataziel:jax_init

+150-63

Commits on Dec 11, 2024

use jaxified logp for initial point evaluation when sampling via Jax
Goose
committed
correcting initial point type hinting
Goose
committed
refactor init_jitter inputs
Goose
committed
revert changes to tests
Goose
committed
removed redundant line from test
Goose
committed

Commits on Dec 12, 2024

correct type annotations related to jaxified logp func
Goose
committed
improved docstrings & type annotations
Goose
committed
fix docstring generation error
Goose
committed
removed type_checking to fix sphinx-autodoc
Goose
committed