diff --git a/.claude/skills/prepare-release/SKILL.md b/.claude/skills/prepare-release/SKILL.md
new file mode 100644
index 000000000..74e387299
--- /dev/null
+++ b/.claude/skills/prepare-release/SKILL.md
@@ -0,0 +1,144 @@
+---
+name: prepare-release
+description: >
+  Automate the HSSM pre-release workflow: update changelog, update docs announcement banner,
+  build docs locally, and create a draft GitHub release. Use this skill whenever the user mentions
+  "prepare release", "release prep", "pre-release checklist", "update changelog for release",
+  or any variation of getting ready to publish a new version of HSSM. Also trigger when the user
+  references the release workflow steps (changelog + announcement + docs build + draft release).
+---
+
+# Prepare Release
+
+This skill automates the pre-release checklist for the HSSM package. It assumes the version
+in `pyproject.toml` has already been bumped. The skill walks through four steps:
+
+1. Update the changelog (`docs/changelog.md`)
+2. Update the announcement banner (`docs/overrides/main.html`)
+3. Build docs locally to verify they compile
+4. Run the notebook check workflow in CI
+5. Create a draft GitHub release
+
+## Conventions
+
+- **Version in pyproject.toml**: no prefix (e.g., `0.3.0`)
+- **Git tags**: `v` prefix (e.g., `v0.3.0`)
+- **Changelog headings**: no prefix (e.g., `### 0.3.0`)
+- **Doc deployment**: happens automatically in CI when a release is *published* — this skill only creates a *draft*
+
+## Step 1: Read the current version
+
+Read `pyproject.toml` and extract the `version` field (line 3). Store this as `VERSION` for the remaining steps.
+
+## Step 2: Gather changes since the last release
+
+Find the most recent git tag:
+
+```bash
+git tag --list 'v*' --sort=-version:refname | head -1
+```
+
+Then collect commits since that tag:
+
+```bash
+git log <LATEST_TAG>..HEAD --oneline
+```
+
+Also look at merged PRs specifically, since they tend to be the most meaningful changelog entries:
+
+```bash
+git log <LATEST_TAG>..HEAD --oneline --grep="Merge pull request"
+```
+
+If useful, also check `gh pr list --state merged --base main --limit 30` for PR titles and descriptions to better understand what each change does.
+
+## Step 3: Draft the changelog entry
+
+Using the gathered commits and PR info, draft a changelog entry that matches the existing
+format in `docs/changelog.md`. The format is:
+
+```markdown
+### {VERSION}
+
+This version includes the following changes:
+
+1. Description of first change.
+2. Description of second change.
+...
+```
+
+Guidelines for writing the changelog:
+- Synthesize commits into human-readable descriptions — don't just paste commit messages
+- Group related commits into single entries
+- Focus on user-facing changes: new features, bug fixes, breaking changes, new tutorials
+- Skip pure CI/infra changes unless they affect users (e.g., new Python version support)
+- Use the style and tone of existing entries as a guide
+
+**Important: Show the draft changelog to the user and ask for their review before writing it to the file.** Use `AskUserQuestion` or just present it inline and wait for confirmation. The user may want to reword entries, add context, or remove items.
+
+## Step 4: Write the changelog
+
+After user approval, insert the new version section into `docs/changelog.md` immediately after the `# Changelog` heading on line 1. Add a blank line before and after the new section.
+
+## Step 5: Update the announcement banner
+
+Edit `docs/overrides/main.html`. Find the line containing the version announcement (pattern: `v{OLD_VERSION} is released!`) and replace it with:
+
+```html
+<span class="right-margin"> v{VERSION} is released! </span>
+```
+
+Only change the version number — leave the surrounding HTML and Jinja2 template structure untouched.
+
+## Step 6: Build docs locally
+
+Run:
+
+```bash
+uv run --group notebook --group docs mkdocs build
+```
+
+Check that the command exits with code 0. If it fails:
+- Show the error output to the user
+- Do NOT proceed to creating the release
+- Help debug the issue
+
+If it succeeds, confirm to the user and move on.
+
+## Step 7: Run the notebook check workflow
+
+Trigger the `check_notebooks.yml` workflow on the current branch to verify all notebooks execute successfully:
+
+```bash
+gh workflow run "Check notebooks" --ref $(git branch --show-current)
+```
+
+Then monitor the run:
+
+```bash
+# Wait a few seconds for the run to register, then find it
+gh run list --workflow="Check notebooks" --limit 1 --json databaseId,status,conclusion
+```
+
+Use `gh run watch <RUN_ID>` to stream the status, or poll with `gh run view <RUN_ID>` periodically.
+
+- If the workflow **succeeds**, confirm to the user and proceed to the draft release.
+- If the workflow **fails**, show the user the failure details using `gh run view <RUN_ID> --log-failed` and help debug. Do NOT proceed to creating the release until notebooks pass.
+
+## Step 8: Create a draft GitHub release
+
+Run:
+
+```bash
+gh release create v{VERSION} --draft --title "v{VERSION}" --generate-notes --target main
+```
+
+The `--draft` flag ensures nothing is published (and therefore no CI release pipeline is triggered).
+The `--generate-notes` flag uses GitHub's auto-generated release notes from merged PRs.
+
+## Step 9: Report
+
+Tell the user:
+- The draft release URL (from the `gh release create` output)
+- A summary of what was done: changelog updated, banner updated, docs build verified, draft release created
+- Remind them that publishing the release will trigger the full CI pipeline (tests, PyPI publish, docs deploy)
diff --git a/.github/workflows/check_notebooks.yml b/.github/workflows/check_notebooks.yml
index 0baff2637..ca704895b 100644
--- a/.github/workflows/check_notebooks.yml
+++ b/.github/workflows/check_notebooks.yml
@@ -34,6 +34,8 @@ jobs:
             "rlssm_rlwm_model.ipynb"
             "rlssm_tutorial.ipynb"
             "add_custom_rlssm_model.ipynb"
+            "hssm_tutorial_workshop_1.ipynb"
+            "hssm_tutorial_workshop_2.ipynb"
           )
 
           EXIT_CODE=0
diff --git a/.github/workflows/prepare-release.yml b/.github/workflows/prepare-release.yml
new file mode 100644
index 000000000..9fda787d4
--- /dev/null
+++ b/.github/workflows/prepare-release.yml
@@ -0,0 +1,150 @@
+name: Prepare Release (AI-assisted)
+
+# This workflow uses Claude to automate pre-release tasks:
+# 1. Generate changelog from merged PRs
+# 2. Update the docs announcement banner
+# 3. Verify docs build
+# 4. Create a draft GitHub release
+#
+# Triggered manually. Claude generates a PR with the changelog and banner
+# updates, then native steps verify the docs build and create the draft release.
+
+on:
+  workflow_dispatch:
+    inputs:
+      version:
+        description: "Version to release (e.g., 0.3.0). Must match the version in pyproject.toml."
+        required: true
+        type: string
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  prepare-changelog-and-banner:
+    name: Generate changelog and update banner via Claude
+    runs-on: ubuntu-latest
+    outputs:
+      pr_number: ${{ steps.create-pr.outputs.pr_number }}
+      branch: ${{ steps.branch.outputs.name }}
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0 # Full history needed for git log
+
+      - name: Create release prep branch
+        id: branch
+        run: |
+          BRANCH="release-prep/v${{ inputs.version }}"
+          git checkout -b "$BRANCH"
+          echo "name=$BRANCH" >> "$GITHUB_OUTPUT"
+
+      - name: Run Claude to update changelog and banner
+        uses: anthropics/claude-code-action@beta
+        with:
+          prompt: |
+            You are preparing release v${{ inputs.version }} of the HSSM package.
+
+            ## Task 1: Update the changelog
+
+            1. Find the latest git tag: `git tag --list 'v*' --sort=-version:refname | head -1`
+            2. Gather changes: `git log <tag>..HEAD --oneline`
+            3. Read `docs/changelog.md` to see the existing format
+            4. Insert a new section after the `# Changelog` heading with this format:
+
+            ```
+            ### ${{ inputs.version }}
+
+            This version includes the following changes:
+
+            1. Human-readable description of change
+            2. ...
+            ```
+
+            Synthesize commits into user-facing descriptions. Focus on features, bug fixes,
+            and breaking changes. Skip CI-only changes.
+
+            ## Task 2: Update the announcement banner
+
+            Edit `docs/overrides/main.html`. Find the line with the version announcement
+            (pattern: `vX.Y.Z is released!`) and replace the version with v${{ inputs.version }}.
+
+            ## Important
+            - Changelog headings use NO "v" prefix (e.g., ### 0.3.0)
+            - Git tags use "v" prefix (e.g., v0.3.0)
+            - Only modify docs/changelog.md and docs/overrides/main.html
+          anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
+
+      - name: Commit changes
+        run: |
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
+          git add docs/changelog.md docs/overrides/main.html
+          git commit -m "docs: update changelog and banner for v${{ inputs.version }}"
+          git push origin "${{ steps.branch.outputs.name }}"
+
+      - name: Create pull request
+        id: create-pr
+        run: |
+          PR_URL=$(gh pr create \
+            --title "docs: prepare release v${{ inputs.version }}" \
+            --body "## Summary
+          - Updated \`docs/changelog.md\` with v${{ inputs.version }} entries (AI-generated from git log)
+          - Updated announcement banner in \`docs/overrides/main.html\`
+
+          ## Review checklist
+          - [ ] Changelog entries are accurate and well-written
+          - [ ] No important changes are missing
+          - [ ] Banner version is correct
+
+          > Auto-generated by the prepare-release workflow using Claude." \
+            --base main \
+            --head "${{ steps.branch.outputs.name }}")
+          echo "pr_number=$(echo $PR_URL | grep -o '[0-9]*$')" >> "$GITHUB_OUTPUT"
+          echo "Created PR: $PR_URL"
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+  verify-docs-build:
+    name: Verify docs build
+    needs: prepare-changelog-and-banner
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout PR branch
+        uses: actions/checkout@v6
+        with:
+          ref: ${{ needs.prepare-changelog-and-banner.outputs.branch }}
+
+      - name: Setup environment
+        uses: ./.github/setup-env
+        with:
+          python-version: "3.13"
+
+      - name: Build docs
+        run: uv run --group notebook --group docs mkdocs build
+
+  create-draft-release:
+    name: Create draft GitHub release
+    needs: [prepare-changelog-and-banner, verify-docs-build]
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - name: Create draft release
+        run: |
+          gh release create "v${{ inputs.version }}" \
+            --draft \
+            --title "v${{ inputs.version }}" \
+            --generate-notes \
+            --target main
+          echo "Draft release created: https://github.com/${{ github.repository }}/releases"
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 1d3e25162..fed323f7f 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -13,6 +13,7 @@ on:
 jobs:
   run_tests:
     runs-on: ubuntu-latest
+    timeout-minutes: 90
     if: ${{ ! contains(github.event.head_commit.message, '[skip fast tests]') }}
     env:
       PYTENSOR_FLAGS: "blas__ldflags=-L/usr/lib/x86_64-linux-gnu -lblas -llapack"
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 000000000..7d8d50ce1
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,108 @@
+# HSSM — Project Context for Claude
+
+## What is HSSM?
+
+HSSM (Hierarchical Sequential Sampling Models) is a Python package for Bayesian inference on sequential sampling models (DDM, LBA, etc.) using PyMC and bambi. It provides a high-level API for defining, fitting, and analyzing these models.
+
+## Project Structure
+
+```
+src/hssm/          # Main package source code
+  hssm.py          # Core HSSM class — the main user-facing API
+  config.py        # Model configuration
+  modelconfig/     # 20+ model config files (DDM, LBA, RDM, etc.)
+  likelihoods/     # Likelihood functions
+  distribution_utils/
+  param/           # Parameter handling
+  plotting/        # Plotting utilities
+  rl/              # Reinforcement learning SSM models
+tests/             # pytest test suite
+docs/              # MkDocs documentation source
+  tutorials/       # Jupyter notebook tutorials (30+ notebooks)
+  changelog.md     # Release changelog
+  overrides/       # MkDocs theme overrides (banner, etc.)
+.github/workflows/ # CI workflows
+.claude/skills/    # Claude Code skills (e.g. prepare-release)
+```
+
+## Build & Tooling
+
+- **Build system:** hatchling
+- **Package manager:** uv (with `uv.lock`)
+- **Python:** >=3.11, <3.14
+- **Linting/formatting:** ruff (via pre-commit)
+- **Type checking:** mypy
+- **Pre-commit hooks:** end-of-file-fixer, trailing-whitespace, ruff, ruff-format, mypy
+
+## Dependency Management
+
+HSSM uses PEP 735 dependency groups alongside traditional optional-dependencies:
+
+- **Core deps** (`[project.dependencies]`): pymc, bambi, numpyro, jax, arviz, etc.
+- **Optional extras** (`[project.optional-dependencies]`): `cuda12`
+- **Dependency groups** (`[dependency-groups]`):
+  - `dev` — pytest, ruff, mypy, pre-commit, coverage
+  - `notebook` — jupyterlab, nbconvert, graphviz, bayesflow (dev branch), keras
+  - `docs` — mkdocs-material, mkdocs-jupyter, mkdocstrings-python
+
+## Common Commands
+
+```bash
+# Install all dev dependencies
+uv sync --group dev --group notebook --group docs
+
+# Run tests
+uv run pytest tests/
+
+# Run slow tests (marked with @pytest.mark.slow)
+uv run pytest tests/ --runslow
+
+# Build docs
+uv run --group notebook --group docs mkdocs build
+
+# Serve docs locally
+uv run --group notebook --group docs mkdocs serve
+
+# Execute a single notebook (for verification)
+uv run --group notebook jupyter nbconvert --ExecutePreprocessor.timeout=10000 --to notebook --execute docs/tutorials/<notebook>.ipynb
+
+# Lint & format
+uv run ruff check .
+uv run ruff format .
+```
+
+## Key Patterns
+
+### HSSM's `**kwargs` passthrough
+
+`HSSM.sample()` passes `**kwargs` through to `bambi.Model.fit()`, which in turn passes them to PyMC's `pm.sample()`. So parameters like `cores`, `chains`, `nuts_sampler`, `target_accept`, etc. are valid even though they don't appear in HSSM's own signature. Similarly, the HSSM constructor passes `**kwargs` to `bambi.Model()`, so bambi parameters like `noncentered` are valid.
+
+### Notebook execution in CI
+
+Two separate skip mechanisms for notebooks:
+1. **mkdocs** `execute_ignore` in `mkdocs.yml` — skips execution during docs build
+2. **CI** `SKIP_NOTEBOOKS` in `.github/workflows/check_notebooks.yml` — skips during notebook CI checks
+
+### Versioning conventions
+
+- `pyproject.toml` version: no prefix (e.g., `0.3.0`)
+- Git tags: `v` prefix (e.g., `v0.3.0`)
+- Changelog headings: no prefix (e.g., `### 0.3.0`)
+
+## CI Workflows
+
+| Workflow | Purpose |
+|----------|---------|
+| `run_tests.yml` | Fast test suite |
+| `run_slow_tests.yml` | Slow tests (`@pytest.mark.slow`) |
+| `linting_and_type_checking.yml` | ruff + mypy |
+| `check_notebooks.yml` | Execute all non-skipped notebooks |
+| `coverage.yml` | Code coverage |
+| `build_docs.yml` | Build documentation |
+| `build_and_publish.yml` | Release to PyPI (triggered on release publish) |
+
+## Known Issues / Notes
+
+- **Multiprocessing in notebooks:** PyMC's NUTS sampler with `cores>1` can cause `EOFError` in notebook execution contexts. Fix by adding `cores=1, chains=1` to `sample()` calls.
+- **bayesflow:** Currently installed from dev branch (`git+https://github.com/bayesflow-org/bayesflow@dev`) in the notebook group only, because PyPI release (2.0.8) doesn't yet include `RatioApproximator`. Not exposed as a user-facing optional extra until a stable release is available.
+- **Notebook kernels:** Some notebooks may have hardcoded kernel specs (e.g., `hssm-dev`). These should use `python3` instead.
diff --git a/docs/changelog.md b/docs/changelog.md
index c72e53e53..2e9c1f0f6 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -1,5 +1,20 @@
 # Changelog
 
+### 0.3.0
+
+This version includes the following changes:
+
+1. Support for **choice-only models**: the HSSM class, data validator, distributions, and model configs now handle models without reaction times, including a softmax likelihood family.
+2. **Racing Diffusion Model (RDM3)**: analytical likelihood, model configuration, and tests for a 3-choice RDM, with safe negative-RT handling and JAX backend support.
+3. **Poisson Race model**: initial implementation of the Poisson race model.
+4. **Hidden Markov Model (HMM-SSM)** example notebook.
+5. **BayesFlow LRE integration** through HSSM for likelihood ratio estimation.
+6. **RLSSM config system**: new `RLSSMConfig` and `BaseModelConfig` classes with comprehensive validation, plus a generalized RL likelihood builder supporting multiple computed parameters.
+7. New tutorials: choice-only modeling, first HMM-SSM example, and updated existing tutorials.
+8. `DataValidator` refactored to `DataValidatorMixin` for improved extensibility.
+9. Bug fixes: default prior assignment, dimensionality errors with Bambi 0.17.0, negative-RT checks on missing data, flaky tests.
+10. Infrastructure: Python 3.13 support, restructured CI test workflows with coverage reporting, updated `model.sample()` API to match Bambi conventions.
+
 ### 0.2.12
 
 This version includes the following changes:
diff --git a/docs/overrides/main.html b/docs/overrides/main.html
index 6023f3195..26da2d2d1 100644
--- a/docs/overrides/main.html
+++ b/docs/overrides/main.html
@@ -5,7 +5,7 @@
   </span>
   Navigate the site here!
 </span>
-<span class="right-margin"> v0.2.12 is released! </span>
+<span class="right-margin"> v0.3.0 is released! </span>
 <span>
   <span class="twemoji">
     {% include ".icons/material/head-question.svg" %}
diff --git a/docs/tutorials/bayesflow_lre_integration.ipynb b/docs/tutorials/bayesflow_lre_integration.ipynb
index 208239c95..5d56640cb 100644
--- a/docs/tutorials/bayesflow_lre_integration.ipynb
+++ b/docs/tutorials/bayesflow_lre_integration.ipynb
@@ -105,14 +105,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "LOAD_PRETRAINED = True\n",
-    "SAVE_DIR = os.path.join(\"data\", \"bayesflow_models\")\n",
-    "PRE_TRAINED_NAME = \"ddm_ratio_approximator.keras\""
-   ]
+   "source": "LOAD_PRETRAINED = True\nSAVE_DIR = os.path.join(\"data\", \"bayesflow_models\")\nPRE_TRAINED_NAME = \"ddm_ratio_approximator.keras\""
   },
   {
    "cell_type": "markdown",
@@ -335,46 +331,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Skipping training (LOAD_PRETRAINED = True).\n"
-     ]
-    }
-   ],
-   "source": [
-    "if not LOAD_PRETRAINED:\n",
-    "    adapter = bf.approximators.RatioApproximator.build_adapter(\n",
-    "        inference_variables=[\"v\", \"a\", \"z\", \"t\"],\n",
-    "        inference_conditions=[\"obs\"],\n",
-    "    )\n",
-    "\n",
-    "    ratio_approximator = bf.approximators.RatioApproximator(\n",
-    "        adapter=adapter,\n",
-    "        classifier_network=bf.networks.MLP(widths=[256, 256, 256]),\n",
-    "        standardize=None,\n",
-    "    )\n",
-    "\n",
-    "    ratio_approximator.compile(optimizer=\"adam\")\n",
-    "    history = ratio_approximator.fit(\n",
-    "        simulator=simulator,\n",
-    "        epochs=150,\n",
-    "        num_batches=200,\n",
-    "        batch_size=64,\n",
-    "    )\n",
-    "\n",
-    "    bf.diagnostics.plots.loss(history)\n",
-    "\n",
-    "    os.makedirs(SAVE_DIR, exist_ok=True)\n",
-    "    ratio_approximator.save(os.path.join(SAVE_DIR, PRE_TRAINED_NAME))\n",
-    "    print(f\"Saved to {SAVE_DIR}/{PRE_TRAINED_NAME}\")\n",
-    "else:\n",
-    "    print(\"Skipping training (LOAD_PRETRAINED = True).\")"
-   ]
+   "outputs": [],
+   "source": "if not LOAD_PRETRAINED:\n    adapter = bf.approximators.RatioApproximator.build_adapter(\n        inference_variables=[\"v\", \"a\", \"z\", \"t\"],\n        inference_conditions=[\"obs\"],\n    )\n\n    ratio_approximator = bf.approximators.RatioApproximator(\n        adapter=adapter,\n        inference_network=bf.networks.MLP(widths=[256, 256, 256]),\n        standardize=None,\n    )\n\n    ratio_approximator.compile(optimizer=\"adam\")\n    history = ratio_approximator.fit(\n        simulator=simulator,\n        epochs=150,\n        num_batches=200,\n        batch_size=64,\n    )\n\n    bf.diagnostics.plots.loss(history)\n\n    os.makedirs(SAVE_DIR, exist_ok=True)\n    ratio_approximator.save(os.path.join(SAVE_DIR, PRE_TRAINED_NAME))\n    print(f\"Saved to {SAVE_DIR}/{PRE_TRAINED_NAME}\")\nelse:\n    print(\"Skipping training (LOAD_PRETRAINED = True).\")"
   },
   {
    "cell_type": "markdown",
@@ -423,48 +383,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "def make_jax_log_ratio_fn(ratio_approximator):\n",
-    "    \"\"\"Extract a pure JAX single-trial log-ratio callable.\n",
-    "\n",
-    "    Returns\n",
-    "    -------\n",
-    "    Callable\n",
-    "        f(obs_i, v, a, z, t) -> scalar log_ratio\n",
-    "        where obs_i is shape (2,) = [rt, choice] and params are scalars.\n",
-    "    \"\"\"\n",
-    "    classifier = ratio_approximator.classifier_network\n",
-    "    projector = ratio_approximator.projector\n",
-    "\n",
-    "    std_layers = getattr(ratio_approximator, \"standardize_layers\", None) or {}\n",
-    "    std_keys = getattr(ratio_approximator, \"standardize\", None) or []\n",
-    "\n",
-    "    def single_trial_log_ratio(obs_i, *params):\n",
-    "        inf_vars = jnp.array([jnp.squeeze(p) for p in params])\n",
-    "        inf_conds = jnp.atleast_1d(obs_i)\n",
-    "\n",
-    "        if \"inference_variables\" in std_keys and \"inference_variables\" in std_layers:\n",
-    "            inf_vars = std_layers[\"inference_variables\"](\n",
-    "                inf_vars[None, :], training=False\n",
-    "            )[0]\n",
-    "        if \"inference_conditions\" in std_keys and \"inference_conditions\" in std_layers:\n",
-    "            inf_conds = std_layers[\"inference_conditions\"](\n",
-    "                inf_conds[None, :], training=False\n",
-    "            )[0]\n",
-    "\n",
-    "        classifier_input = jnp.concatenate([inf_vars, inf_conds])\n",
-    "        hidden = classifier(classifier_input[None, :], training=False)\n",
-    "        logits = projector(hidden, training=False)\n",
-    "        return jnp.squeeze(logits).astype(jnp.float64)\n",
-    "\n",
-    "    return single_trial_log_ratio\n",
-    "\n",
-    "\n",
-    "single_trial_fn = make_jax_log_ratio_fn(ratio_approximator)"
-   ]
+   "source": "def make_jax_log_ratio_fn(ratio_approximator):\n    \"\"\"Extract a pure JAX single-trial log-ratio callable.\n\n    Returns\n    -------\n    Callable\n        f(obs_i, v, a, z, t) -> scalar log_ratio\n        where obs_i is shape (2,) = [rt, choice] and params are scalars.\n    \"\"\"\n    inference_net = ratio_approximator.inference_network\n    projector = ratio_approximator.projector\n\n    std_layers = getattr(ratio_approximator, \"standardize_layers\", None) or {}\n    std_keys = getattr(ratio_approximator, \"standardize\", None) or []\n\n    def single_trial_log_ratio(obs_i, *params):\n        inf_vars = jnp.array([jnp.squeeze(p) for p in params])\n        inf_conds = jnp.atleast_1d(obs_i)\n\n        if \"inference_variables\" in std_keys and \"inference_variables\" in std_layers:\n            inf_vars = std_layers[\"inference_variables\"](\n                inf_vars[None, :], training=False\n            )[0]\n        if \"inference_conditions\" in std_keys and \"inference_conditions\" in std_layers:\n            inf_conds = std_layers[\"inference_conditions\"](\n                inf_conds[None, :], training=False\n            )[0]\n\n        classifier_input = jnp.concatenate([inf_vars, inf_conds])\n        hidden = inference_net(classifier_input[None, :], training=False)\n        logits = projector(hidden, training=False)\n        return jnp.squeeze(logits).astype(jnp.float64)\n\n    return single_trial_log_ratio\n\n\nsingle_trial_fn = make_jax_log_ratio_fn(ratio_approximator)"
   },
   {
    "cell_type": "code",
diff --git a/docs/tutorials/data/bayesflow_models/ddm_ratio_approximator.keras b/docs/tutorials/data/bayesflow_models/ddm_ratio_approximator.keras
index bf191d5cf..bba6044a1 100644
Binary files a/docs/tutorials/data/bayesflow_models/ddm_ratio_approximator.keras and b/docs/tutorials/data/bayesflow_models/ddm_ratio_approximator.keras differ
diff --git a/docs/tutorials/initial_values.ipynb b/docs/tutorials/initial_values.ipynb
index f7b096d52..63090c8d4 100644
--- a/docs/tutorials/initial_values.ipynb
+++ b/docs/tutorials/initial_values.ipynb
@@ -148,57 +148,8 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Initializing NUTS using adapt_diag...\n",
-      "Multiprocess sampling (4 chains in 4 jobs)\n",
-      "NUTS: [a, t, z, v]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd4dadcf76ea439c8ae1347f223e60ee",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Output()"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
-      ],
-      "text/plain": []
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Sampling 4 chains for 100 tune and 100 draw iterations (400 + 400 draws total) took 33 seconds.\n",
-      "The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details\n",
-      "The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details\n",
-      "100%|██████████| 400/400 [00:11<00:00, 34.16it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "idata = model.sample(draws=100,\n",
-    "                     tune=100,\n",
-    "                     sampler=\"pymc\",\n",
-    "                     initvals = my_initvals)"
-   ]
+   "outputs": [],
+   "source": "idata = model.sample(draws=100,\n                     tune=100,\n                     sampler=\"pymc\",\n                     cores=1,\n                     chains=1,\n                     initvals = my_initvals)"
   },
   {
    "cell_type": "code",
@@ -913,67 +864,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using default initvals. \n",
-      "\n",
-      "The model has already been sampled. Overwriting the previous inference object. Any previous reference to the inference object will still point to the old object.\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Only 10 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.\n",
-      "Initializing NUTS using adapt_diag...\n",
-      "Multiprocess sampling (4 chains in 4 jobs)\n",
-      "NUTS: [a, t, z, v]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6074108665f141eda1ef0d57731edaf5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Output()"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
-      ],
-      "text/plain": []
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Sampling 4 chains for 1 tune and 10 draw iterations (4 + 40 draws total) took 1 seconds.\n",
-      "There were 40 divergences after tuning. Increase `target_accept` or reparameterize.\n",
-      "The number of samples is too small to check convergence reliably.\n",
-      "100%|██████████| 40/40 [00:00<00:00, 54.07it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "model._initvals[\"z\"] = np.array(0.6, dtype=\"float32\")\n",
-    "idata2 = model.sample(draws=10, tune=1)"
-   ]
+   "outputs": [],
+   "source": "model._initvals[\"z\"] = np.array(0.6, dtype=\"float32\")\nidata2 = model.sample(draws=10, tune=1, cores=1, chains=1)"
   },
   {
    "cell_type": "code",
diff --git a/docs/tutorials/plotting.ipynb b/docs/tutorials/plotting.ipynb
index a09a602ae..7499f38e8 100644
--- a/docs/tutorials/plotting.ipynb
+++ b/docs/tutorials/plotting.ipynb
@@ -3946,41 +3946,8 @@
    "execution_count": null,
    "id": "1ed14656",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using default initvals. \n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.11/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.\n",
-      "  pmap_numpyro = MCMC(\n",
-      "sample: 100%|██████████| 1000/1000 [02:00<00:00,  8.29it/s, 31 steps of size 3.26e-02. acc. prob=0.90]\n",
-      "sample: 100%|██████████| 1000/1000 [02:23<00:00,  6.99it/s, 51 steps of size 3.07e-02. acc. prob=0.92]\n",
-      "There were 988 divergences after tuning. Increase `target_accept` or reparameterize.\n",
-      "We recommend running at least 4 chains for robust computation of convergence diagnostics\n",
-      "The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details\n",
-      "The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details\n"
-     ]
-    }
-   ],
-   "source": [
-    "idata_race = race_model.sample(\n",
-    "    sampler=\"numpyro\"\n",
-    "    chains=2,\n",
-    "    cores=2,\n",
-    "    chain_method=\"vectorized\",\n",
-    "    draws=500,\n",
-    "    tune=500,\n",
-    "    idata_kwargs=dict(log_likelihood=False),  # no need to return likelihoods here\n",
-    ")"
-   ]
+   "outputs": [],
+   "source": "idata_race = race_model.sample(\n    sampler=\"numpyro\",\n    chains=2,\n    cores=2,\n    chain_method=\"vectorized\",\n    draws=500,\n    tune=500,\n    idata_kwargs=dict(log_likelihood=False),  # no need to return likelihoods here\n)"
   },
   {
    "cell_type": "markdown",
diff --git a/docs/tutorials/poisson_race.ipynb b/docs/tutorials/poisson_race.ipynb
index fcc55dc22..9b841bb34 100644
--- a/docs/tutorials/poisson_race.ipynb
+++ b/docs/tutorials/poisson_race.ipynb
@@ -246,7 +246,7 @@
       "  res = dot(at, bt)\n",
       "/Users/hayden/miniconda3/envs/hssm-dev/lib/python3.12/site-packages/numpy/_core/numeric.py:1211: RuntimeWarning: invalid value encountered in dot\n",
       "  res = dot(at, bt)\n",
-      "100%|██████████| 6000/6000 [00:00<00:00, 18880.52it/s]\n"
+      "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 6000/6000 [00:00<00:00, 18880.52it/s]\n"
      ]
     }
    ],
@@ -472,9 +472,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "hssm-dev",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "hssm-dev"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
diff --git a/mkdocs.yml b/mkdocs.yml
index a2e04bacf..93b2696fd 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -103,13 +103,14 @@ plugins:
         - tutorials/rlssm_rlwm_model.ipynb
         - tutorials/tutorial_bayeux.ipynb
         - tutorials/bayesflow_lre_integration.ipynb
+        - tutorials/poisson_race.ipynb
         - .ipynb_checkpoints/*.ipynb
       allow_errors: false
   - mkdocstrings:
       default_handler: python
       handlers:
         python:
-          import:
+          inventories:
             - https://docs.python.org/3/objects.inv
             - https://mkdocstrings.github.io/objects.inv
             - https://mkdocstrings.github.io/griffe/objects.inv
diff --git a/pyproject.toml b/pyproject.toml
index 996f7ac35..031feee3f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "HSSM"
-version = "0.2.12"
+version = "0.3.0"
 description = "Bayesian inference for hierarchical sequential sampling models."
 authors = [
     { name = "Alexander Fengler", email = "alexander_fengler@brown.edu" },
@@ -36,7 +36,7 @@ dependencies = [
     "pandas>=2.2,<3",
     "pymc>=5.26.0",
     "seaborn>=0.13.2",
-    "ssm-simulators>=0.12.0",
+    "ssm-simulators>=0.12.2",
     "tqdm>=4.66.0",
 ]
 
@@ -47,7 +47,6 @@ repository = "https://github.com/lnccbrown/HSSM"
 
 [project.optional-dependencies]
 cuda12 = ["jax[cuda12]>=0.5.2"]
-bayesflow = ["bayesflow", "keras>=3"]
 
 [dependency-groups]
 dev = [
@@ -58,6 +57,7 @@ dev = [
     "pytest>=8.3.1",
     "pytest-random-order>=1.1.1",
     "pytest-rerunfailures>=15.0",
+    "pytest-timeout>=2.3.1",
     "ruff>=0.15.0",
     "pre-commit>=4.1.0",
     "onnxruntime>=1.17.1",
@@ -75,8 +75,10 @@ notebook = [
     "nbval>=0.11.0",
     "ptpython>=3.0.29",
     "pyarrow>=20.0.0",
-    "lanfactory>=0.5.3",
+    "lanfactory>=0.6.1",
     "zeus-mcmc>=2.5.4",
+    "bayesflow>=2.0.10",
+    "keras>=3",
 ]
 
 docs = [
@@ -215,6 +217,9 @@ addopts = [
     "--cov-report=xml",
     "--exitfirst",
     "--capture=no",
+    "--timeout=360",
+    "--reruns=2",
+    "--reruns-delay=5",
 ]
 markers = [
     "slow: marks tests as slow (add --runslow option to run these tests)",
diff --git a/tests/conftest.py b/tests/conftest.py
index 728504623..878c5e780 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -37,18 +37,18 @@ def data_angle():
 def data_ddm_reg():
     # Generate some fake simulation data
     intercept = 1.5
-    x = np.random.uniform(-0.5, 0.5, size=1000)
-    y = np.random.uniform(-0.5, 0.5, size=1000)
+    x = np.random.uniform(-0.5, 0.5, size=250)
+    y = np.random.uniform(-0.5, 0.5, size=250)
 
     v = intercept + 0.8 * x + 0.3 * y
     true_values = np.column_stack(
-        [v, np.repeat([[1.5, 0.5, 0.5]], axis=0, repeats=1000)]
+        [v, np.repeat([[1.5, 0.5, 0.5]], axis=0, repeats=250)]
     )
 
     dataset_reg_v = hssm.simulate_data(
         model="ddm",
         theta=true_values,
-        size=1,  # Generate one data point for each of the 1000 set of true values
+        size=1,
     )
 
     dataset_reg_v["x"] = x