Example: BayesFlow Likelihood as serialized JAX function
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 124
- Forks
- 24
- Avg merge
- 19h 32m
- Merged PRs (30d)
- 60
Description
BayesFlow uses Keras throughout, which has become backend-agnostic in recent years.
Two key aspects for HSSM:
- Keras supports serialization of models (save/load)
- We should be able to use a Keras model as a simple JAX function finally
To close this issue, produce a simple example of inter-operability between BayesFlow and HSSM.
1. Use `ssm-simulators` to instantiate a basic DDM simulators
2. Train a likelihood networks via BayesFlow
3. Save the model
4. Load this model and turn it into a differentiable likelihood that works with the HSSM low-level interface
5. Translate the use case to the high-level interface
The output here should be one or two Jupyter notebooks that can be part of the docs as initial examples.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the existing HSSM documentation and the low-level and high-level interfaces, then examine how ssm-simulators and BayesFlow are used together. Build one or two Jupyter notebooks covering DDM simulation, likelihood-network training, model save/load, differentiable JAX conversion, and both interfaces. Done means the notebooks run as documentation examples and demonstrate the complete interoperability workflow.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- jupyter-notebook, keras, python
- Domain
- documentation, machine-learning
- Issue type
- Documentation
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100