fix: Colab bootstrap installs no lazily-imported dependency (corner, optax, xxhash, blackjax)

[Jammy2211]

Overview

HowToFit chapter 1 tutorial 5 dies on Colab with ModuleNotFoundError: No module named 'corner' — after the search has run to completion, in the results update that follows it. The Colab bootstrap installs the stack with --no-deps, so any dependency Colab does not itself preinstall must be named by hand in _SHARED_EXTRAS, and corner==2.2.2 — a base dependency of autofit — is not there.

This is the third report of the same defect (8336939 / #166 closed it for emcee and dynesty; a sibling prompt has it open for blackjax). Because each patch needs its own autonerves PyPI release to reach a single user, this task audits the whole list and closes every hole at once, rather than adding one package and waiting for the next report.

Plan

• Add the four lazily-imported dependencies missing from _SHARED_EXTRAS: corner, optax, xxhash, blackjax.
• Correct the two specifiers that have drifted from the PyAutoFit/pyproject.toml the list's own comment claims to track: anesthetic and nautilus-sampler.
• Document the lazy-import trap in the comment above the list, so the next person understands why these are invisible to import autofit and to smoke.
• Extend the existing name-only regression test, and add a specifier-aware test that derives expectations from autofit's pyproject.toml rather than restating literals — restating them is what allowed the drift.
• Absorb the _SHARED_EXTRAS legs of two open sibling prompts so one release closes all of it.

Detailed implementation plan

Affected Repositories

• PyAutoNerves (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoNerves	main	clean (0 changed, 0 behind origin)

Suggested branch: feature/colab-bootstrap-lazy-deps

Root cause

autonerves/setup_colab.py:126 installs with --no-deps:

subprocess.check_call([sys.executable, "-m", "pip", "install", *packages, "--no-deps"])

corner is imported inside corner_cornerpy (PyAutoFit/autofit/non_linear/plot/samples_plotters.py:95), not at module scope. So import autofit succeeds, the model builds, the search runs, and only the post-fit plot detonates. Workspace smoke runs at PYAUTO_TEST_MODE=2, which bypasses the sampler entirely, so no gate has ever executed that line.

Audit: every hole in _SHARED_EXTRAS

Missing entirely:

Package	autofit declares	Imported at	Reaches users via
corner	base, corner==2.2.2	plot/samples_plotters.py:95	every MCMC/nested corner plot — the reported failure
optax	base, optax>=0.2.5 (marker-gated)	search/mle/multi_start_gradient/search.py:898	MultiStartAdam / MultiStartProdigy — HowToFit 6, 7
xxhash	base, xxhash<=3.4.1	graphical/factor_graphs/factor.py:361	graphical/EP factor hashing — HowToFit chapter_advanced
blackjax	optional, blackjax>=1.6.2	BlackJAXNUTS	HowToFit 6, 7

Drifted specifiers. git blame settles these: both were written in 756f98e (2026-07-09) as snapshots of what autofit declared then, and autofit has since moved — drift, not deliberate holds:

Package	list says	autofit declares	Effect under --no-deps
anesthetic	==2.8.14	>=2.9.0 (base)	pins below autofit's own floor
nautilus-sampler	==1.0.4	==1.0.5 (optional)	separate open prompt

Implementation Steps

1. autonerves/setup_colab.py — extend the comment above _SHARED_EXTRAS (:30-37) to name the lazy-import trap.
2. autonerves/setup_colab.py — the list becomes:

_SHARED_EXTRAS = [
    "pyvis==0.3.2",
    "dill==0.4.0",
    "jaxnnls",
    "nautilus-sampler==1.0.5",   # was ==1.0.4 — drifted from autofit's optional extra
    "timeout_decorator==0.5.0",
    "anesthetic>=2.9.0",         # was ==2.8.14 — pinned below autofit's declared floor
    "emcee>=3.1.6",
    "dynesty==2.1.5",
    "corner==2.2.2",             # new — the reported failure
    "optax>=0.2.5",              # new — unmarked; autofit's marker guards Intel macOS, Colab is linux
    "xxhash<=3.4.1",             # new
    "blackjax>=1.6.2",           # new
]

Two notes belong in the code comments because they are not obvious:

• optax's marker in autofit (sys_platform != "darwin" or platform_machine == "arm64") guards Intel macOS, which has no jax wheels. Colab is always linux, so this entry is deliberately unmarked.
• xxhash<=3.4.1 is an upper bound: if Colab ships newer, this downgrades it. That matches autofit's own constraint so it is intended, but it is the one entry that changes an already-working Colab package rather than adding a missing one.

3. test_autonerves/test_setup_colab.py — extend test_every_project_installs_every_sampler (:93-105) required set with the four new packages. Keep its name-only matching: it guards the "missing entirely" class.
4. test_autonerves/test_setup_colab.py — add a specifier-aware test closing the drift class: parse PyAutoFit/pyproject.toml ([project] dependencies + the optional extra) with tomllib; for every _SHARED_EXTRAS entry whose name autofit declares, assert the specifier matches, derived from autofit rather than restated. Skip the environment marker when comparing (optax); skip entries autofit does not declare (jaxnnls); locate PyAutoFit via the sibling-checkout path and pytest.skip cleanly when absent.

Key Files

• autonerves/setup_colab.py — _SHARED_EXTRAS (:38-47), its comment (:30-37), the --no-deps install (:126)
• test_autonerves/test_setup_colab.py — test_every_project_installs_every_sampler (:93-105)

Out of scope (settled with the user)

• Do not change what PyAutoHands injects as the notebook setup cell.
• Do not regenerate or push notebooks in any of the six notebook repos. The setup cell runs pip install autonerves --no-deps unpinned at run time, so a released autonerves fixes every published notebook retroactively, including already-tagged ones. Regeneration carries nothing.
• Do not refactor _SHARED_EXTRAS into a fully metadata-derived install list (offered and declined). The specifier-aware test gets most of the safety at a fraction of the risk.

Sibling prompts absorbed

• draft/bug/autonerves/nautilus_sampler_pin_drifts_from_pyautofit.md — fully absorbed; retire at close-out.
• draft/bug/howtofit/tutorials_6_7_blackjax_never_installed.md — steps 1 and 3 absorbed; re-scope down to its HowToFit/requirements.txt leg (different repo, different install route) rather than closing.

Verification

1. pytest test_autonerves/test_setup_colab.py -q — both tests pass.
2. Revert one specifier (e.g. nautilus to 1.0.4) and confirm the new test fails. A drift test that cannot fail is worth nothing.
3. PyAutoHeart verify_install check F simulates the Colab bootstrap and is the only gate positioned to execute these imports for real.
4. Post-release: a Colab runtime running setup_colab.setup("howtofit") then tutorial 5 to completion, writing its corner plot.

Shipping

Merging fixes nothing on its own — an autonerves PyPI release is what ships this. Latest on PyPI is 2026.9.15.1 (released 2026-09-15). Vitals currently reads RELEASE VALIDATION incomplete — no rehearsal evidence, so confirm the overnight release run has settled before cutting 2026.9.15.2.

Original Prompt

Click to expand starting prompt

Colab bootstrap installs no lazily-imported dependency — corner kills every MCMC fit mid-run

Type: bug
Target: PyAutoNerves
Repos:

• PyAutoNerves
  Themes:
• colab
• dependencies
• tutorials
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Status: formalised
  Consequence: judge
  Review-minutes: 10
  Unattended: ready
  Witness: on a Colab runtime bootstrapped with setup_colab.setup("howtofit"), HowToFit/scripts/chapter_1_introduction/tutorial_5_results_and_samples.py runs to completion and writes its corner plot instead of raising ModuleNotFoundError: No module named 'corner' at the end-of-fit results update.
  Filed: 2026-09-15

Original request (verbatim)

tutorial 5 still gives this

2026-09-15 06:59:25,441 - autofit.non_linear.search.updater - INFO - Creating latent samples by drawing 100 from the PDF.
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_3678/3237708700.py in <cell line: 0>()
    111 )
    112
--> 113 result = search.fit(model=model, analysis=analysis)

8 frames
/usr/local/lib/python3.13/dist-packages/autofit/non_linear/plot/samples_plotters.py in corner_cornerpy(samples, path, filename, format, **kwargs)
     93     pylab.rcParams.update(params)
     94
---> 95     import corner
     96     from corner.core import hist2d as corner_hist2d
     97

ModuleNotFoundError: No module named 'corner'

is there a way to fix it without a whole release?

Reported against HowToFit chapter 1 tutorial 5 on Colab, Python 3.13, autofit
installed from PyPI. The search itself completed (2000 emcee steps, 57 s); the
crash is in the results update that follows it.

Root cause

autonerves/setup_colab.py:126 installs the whole stack with --no-deps:

subprocess.check_call([sys.executable, "-m", "pip", "install", *packages, "--no-deps"])

so any dependency Colab does not itself preinstall has to be named by hand in
_SHARED_EXTRAS (:38-47) or it never lands. corner==2.2.2 is a base
dependency of autofit (PyAutoFit/pyproject.toml:41), is absent from
_SHARED_EXTRAS, and is absent from Colab.

The reason it went unnoticed for so long is the shape of the import. corner is
imported inside corner_cornerpy (autofit/non_linear/plot/samples_plotters.py:95),
not at module scope, so import autofit succeeds, the model builds, the search
runs to completion, and the failure only lands in the results update after the
fit. Nothing before that point can see it.

This is the third instance of the same defect. PyAutoNerves commit 8336939
("fix: install emcee and dynesty in the Colab bootstrap", #166) closed it for
two packages; draft/bug/howtofit/tutorials_6_7_blackjax_never_installed.md
has it open for a third. The list is being patched one report at a time.

Every hole in _SHARED_EXTRAS, found by auditing the whole list

Lazily-imported dependencies absent from the list:

Package	Declared in autofit as	Imported at	Reaches the user via
corner	base, corner==2.2.2	plot/samples_plotters.py:95	every MCMC / nested search's corner plot — the reported failure
optax	base, optax>=0.2.5 (marker-gated)	search/mle/multi_start_gradient/search.py:898	MultiStartAdam / MultiStartProdigy — HowToFit tutorials 6, 7
xxhash	base, xxhash<=3.4.1	graphical/factor_graphs/factor.py:361	graphical / EP factor hashing — HowToFit chapter_advanced
blackjax	optional extra, blackjax>=1.6.2	BlackJAXNUTS	HowToFit tutorials 6, 7 — see the sibling prompt below

Specifiers in the list that have drifted from the file they claim to track
(the comment above _SHARED_EXTRAS says outright "their specifiers track that
file"):

Package	_SHARED_EXTRAS says	autofit declares	Effect under --no-deps
anesthetic	==2.8.14	>=2.9.0 (base)	pins below autofit's own floor
nautilus-sampler	==1.0.4	==1.0.5 (optional)	separate open prompt

optax's marker in autofit (sys_platform != "darwin" or platform_machine == "arm64")
exists because Intel macOS has no jax wheels. Colab is always linux, so the entry
here is unmarked.

xxhash<=3.4.1 is an upper bound: if Colab ships a newer xxhash, adding this
entry downgrades it. That matches autofit's own declared constraint, so it is
the intended behaviour, but it should be a conscious choice rather than a
side-effect. The >=-style entries are no-ops when Colab already satisfies them.

Why no gate sees any of this

Workspace smoke runs at PYAUTO_TEST_MODE=2, which bypasses the sampler, so the
searches are never constructed and the lazy imports are never executed. The
scripts pass green while being unrunnable by a reader on Colab. PyAutoHeart's
verify_install check F simulates the Colab bootstrap and is the only gate in
the right place to catch this.

Sibling prompts covering the same ten lines

• draft/bug/howtofit/tutorials_6_7_blackjax_never_installed.md (Priority: high,
  filed 2026-09-15) — its fix step 1 is blackjax>=1.6.2 into _SHARED_EXTRAS,
  step 3 is the nautilus pin.
• draft/bug/autonerves/nautilus_sampler_pin_drifts_from_pyautofit.md — the
  nautilus drift, and it already asks whether "the other _SHARED_EXTRAS pins
  deserve the same treatment" and recommends deriving the expected specifier from
  PyAutoFit's pyproject.toml rather than restating literals.

Each of these needs its own autonerves PyPI release to reach a single user.
Landing them separately means three releases to close one defect. The plan
should decide whether this prompt absorbs the _SHARED_EXTRAS legs of both
(leaving the blackjax prompt its HowToFit/requirements.txt leg, which is a
different repo and a different install route).

Fix

1. Close the holes in _SHARED_EXTRAS, specifiers tracking
   PyAutoFit/pyproject.toml as the existing entries claim to.
2. Correct the two drifted specifiers.
3. Extend the comment above the list to name the lazy-import trap explicitly —
   a dependency imported inside a function survives import autofit and only
   detonates mid-run, which is why these are missed.
4. Regression test asserting each package lands in the install list. The existing
   test: assert every Colab project installs all three samplers #166 test matches on requirement name with the specifier split off, which is
   why it cannot see a drifted pin; a specifier-aware check that reads
   PyAutoFit's pyproject.toml rather than restating literals would close both
   classes at once.

Out of scope (decided with the user, 2026-09-15)

• Do not change what PyAutoHands injects as the notebook setup cell.
• Do not regenerate or push notebooks in any of the six notebook repos. The
  setup cell runs pip install autonerves --no-deps unpinned at run time, so
  a released autonerves fixes every published notebook retroactively, including
  already-tagged ones. Notebook regeneration carries nothing.
• Do not refactor _SHARED_EXTRAS into a fully metadata-derived install list
  (offered and declined). If it is worth doing, file it as its own prompt.

Shipping

Merging changes nothing on its own — an autonerves PyPI release is what ships
this. Latest on PyPI is 2026.9.15.1, released 2026-09-15. Confirm the overnight
release run has settled before cutting one.

[Jammy2211]

Library PR Created: #169. Second pass adds jax_zero_contour and zeus-mcmc (found by PyAutoLabs/PyAutoHeart#228's Colab gate) plus a PR-time colab-gate.yml workflow; check F reports F|PASS with this branch overlaid. The colab-gate.yml check on this PR goes green only after PyAutoHeart#228 merges.

[Jammy2211]

Shipped

Merged as #169 (merge commit 8eca4b3e8cdf72ed4c13850ed5b333ec7306524d).

The Colab bootstrap now installs every lazily-imported dependency: _SHARED_EXTRAS gains corner, optax, xxhash, blackjax, plus jax_zero_contour and zeus-mcmc found by PyAutoHeart#228's Colab gate; nautilus-sampler and anesthetic specifiers corrected. TestSpecifiersTrackAutofit derives expectations from PyAutoFit's and PyAutoGalaxy's pyproject.toml at run time (185 tests), and the new colab-gate.yml workflow runs Heart check F against the PR checkout (F|PASS on #169).

Release consequence: merging fixes nothing on its own — the published notebooks pip-install autonerves unpinned at run time, so the PyPI release is what reaches them. release-integrate stays RED until the autonerves release carrying PyAutoNerves#169 is on PyPI.

Record: PyAutoMind/complete/2026/09/colab-bootstrap-lazy-deps.md.