From f5d8317ce7391d931637ca5ae2d0788769ec07e5 Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Tue, 11 Aug 2026 19:07:50 -0400
Subject: [PATCH 01/12] Add the end-to-end runner, chest X-ray as a third
 modality, and run provenance

Chest X-ray joins notes and laboratory values through the unified embedding.
CXRMultimodalMIMIC4 adds cxr_only, cxr_labs and cxr_notes_labs. An image event
has a real position on the timeline: StudyDate and StudyTime give hours from
admission, the same convention that laboratory values use. Only a study inside
the observation window enters a sample.

Two fallbacks changed the measurement in silence. _split_dataset changed from
split_by_patient to split_by_sample when the patient split was empty, which
leaks, because the admissions of one patient can then be in both splits. The
reported predictions came from test_loader or val_loader or train_loader, so a
run with no test split reported validation or training performance as test
performance. Both now warn and record what they used.

metrics_history.json records the score of a run but not its conditions, so a
frozen-encoder run and a fine-tuned run are indistinguishable once the job
output is gone. write_run_config records the RESOLVED settings, because
--freeze-encoder is an alias and the raw flag describes the run incorrectly.
Code identity is a git commit and a SHA-256 digest of the package source, since
a cluster run starts from an unpacked archive where git gives no result.

encoder_lr gives a pretrained text encoder a gentler rate than the randomly
initialised layers around it. With the encoder frozen, every encoders.*
parameter has requires_grad=False, so the projection is the only trainable text
parameter and joins the group; otherwise it keeps the base rate.

An epoch mean cannot show the difference between a run that starts badly and a
run that becomes worse inside the epoch, so each epoch also records
train_loss_first_step, train_loss_first100 and train_loss_last100.
---
 .../unified_embedding_e2e_mimic4.py           | 614 +++++++++++++++++-
 pyhealth/datasets/configs/mimic4_cxr.yaml     |   2 +-
 pyhealth/datasets/mimic4.py                   |  13 +-
 pyhealth/tasks/__init__.py                    |   1 +
 pyhealth/tasks/multimodal_mimic4.py           | 197 ++++++
 pyhealth/trainer.py                           |  70 +-
 pyhealth/utils.py                             |  78 +++
 tests/test_run_provenance_and_pathways.py     | 215 ++++++
 8 files changed, 1154 insertions(+), 36 deletions(-)
 create mode 100644 tests/test_run_provenance_and_pathways.py

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 2b4a0dc30..186b2998a 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -10,10 +10,6 @@
     MortalityPredictionStageNetMIMIC4: ICD codes + 10-dim lab vectors,
     patient-level samples aggregated across all admissions.
 
---task icd_labs
-    ICDLabsMIMIC4: ICD codes + 10-dim lab vectors via the unified
-    multimodal pipeline.  No notes required.
-
 --task clinical_notes_icd_labs
     ClinicalNotesICDLabsMIMIC4: discharge/radiology notes + ICD + labs.
     Requires --note-root.  Legacy; ICD codes are discharge-coded (leakage).
@@ -55,6 +51,7 @@
 
 import argparse
 import csv
+import warnings
 from pathlib import Path
 from typing import Any, Tuple
 
@@ -71,17 +68,21 @@
     split_by_sample,
 )
 from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel
+from pyhealth.models.embedding import VisionEmbeddingModel
 from pyhealth.models.bottleneck_transformer import BottleneckTransformer
 from pyhealth.models.ehrmamba import EHRMamba
 from pyhealth.models.jamba_ehr import JambaEHR
+from pyhealth.processors import fit_lab_standardizer, lab_standardizer_fit_scope
 from pyhealth.tasks import MortalityPredictionStageNetMIMIC4
 from pyhealth.tasks.multimodal_mimic4 import (
     ClinicalNotesICDLabsMIMIC4,
     ICDLabsMIMIC4,
+    LabsOnlyMIMIC4,
     NotesLabsMIMIC4,
+    CXRMultimodalMIMIC4,
 )
 from pyhealth.trainer import Trainer
-from pyhealth.utils import set_seed
+from pyhealth.utils import set_seed, write_run_config
 
 
 def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
@@ -96,11 +97,44 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
     if args.task == "icd_labs":
         ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"]
 
+    if args.task in ("notes_labs", "notes_only"):
+        if not args.note_root:
+            raise ValueError(f"--task {args.task} requires --note-root.")
+        note_tables = [getattr(args, "note_source", "discharge")]
+        # Load ICD tables only when explicitly requested (they are discharge-coded).
+        ehr_tables = (
+            ["diagnoses_icd", "procedures_icd", "labevents"]
+            if args.icd_codes
+            else ["labevents"]
+        )
+        if args.include_vitals:
+            if "chartevents" not in ehr_tables:
+                ehr_tables.append("chartevents")
+
+    if args.task == "labs_only":
+        # Pure EHR baseline: only labevents, no notes, no ICD codes.
+        ehr_tables = ["labevents"]
+        note_tables = None
+
+    cxr_tables = None
+    if args.task in ("cxr_only", "cxr_labs", "cxr_notes_labs"):
+        if not args.cxr_root:
+            raise ValueError(f"--task {args.task} requires --cxr-root.")
+        # ``metadata`` supplies image_path, StudyDate/StudyTime, and ViewPosition.
+        cxr_tables = ["metadata"]
+        ehr_tables = ["labevents"] if args.task != "cxr_only" else []
+        if args.task == "cxr_notes_labs":
+            if not args.note_root:
+                raise ValueError("--task cxr_notes_labs requires --note-root.")
+            note_tables = [getattr(args, "note_source", "discharge")]
+
     return MIMIC4Dataset(
         ehr_root=args.ehr_root,
         ehr_tables=ehr_tables,
         note_root=args.note_root if note_tables else None,
         note_tables=note_tables,
+        cxr_root=args.cxr_root if cxr_tables else None,
+        cxr_tables=cxr_tables,
         cache_dir=args.cache_dir,
         dev=args.dev if args.dev else False,
         num_workers=args.num_workers,
@@ -114,27 +148,120 @@ def _build_task(args: argparse.Namespace):
         return ICDLabsMIMIC4(window_hours=args.observation_window_hours)
     if args.task == "clinical_notes_icd_labs":
         return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours)
-    if args.task == "notes_labs":
-        return NotesLabsMIMIC4(
+    if args.task in ("notes_labs", "notes_only"):
+        task_kwargs = dict(
             window_hours=args.observation_window_hours,
             include_icd=args.icd_codes,
             include_vitals=args.include_vitals,
+            include_labs=(args.task != "notes_only"),
+            note_extraction=getattr(args, "note_extraction", "regex"),
+            note_source=getattr(args, "note_source", "discharge"),
+            discharge_note_policy=getattr(
+                args, "discharge_note_policy", "extraction"),
+        )
+        # Only pass text_normalize when actually requested, so this script still
+        # runs against a checkout whose task class predates the parameter.
+        _tn = getattr(args, "text_normalize", "none")
+        if _tn and _tn != "none":
+            task_kwargs["text_normalize"] = _tn
+        task = NotesLabsMIMIC4(**task_kwargs)
+        if args.tokenizer_model:
+            schema_key = "admission_note_times"
+            _, opts = task.input_schema[schema_key]
+            task.input_schema[schema_key] = ("tuple_time_text", {**opts, "tokenizer_model": args.tokenizer_model})
+            print(f"[tokenizer] Overriding tokenizer_model → {args.tokenizer_model}")
+        # Note token budget. The processor default (128) truncates ~96% of
+        # extracted discharge notes, so the encoder sees only their first ~20%.
+        # input_schema is part of the task cache key, so each budget caches apart.
+        _nml = getattr(args, "note_max_length", None)
+        if _nml:
+            # Fail loudly rather than let HF silently clamp to the encoder's
+            # position-embedding limit (BERT/Bio_ClinicalBERT = 512).
+            _tokmodel = args.tokenizer_model or task.input_schema[
+                "admission_note_times"][1].get("tokenizer_model", "")
+            if _nml > 512 and "longformer" not in _tokmodel.lower():
+                raise SystemExit(
+                    f"--note-max-length {_nml} exceeds the 512 position-embedding "
+                    f"limit of {_tokmodel!r}. Use --tokenizer-model "
+                    f"yikuan8/Clinical-Longformer for budgets >512."
+                )
+            schema_key = "admission_note_times"
+            _, opts = task.input_schema[schema_key]
+            task.input_schema[schema_key] = (
+                "tuple_time_text", {**opts, "max_length": _nml}
+            )
+            print(f"[tokenizer] note max_length → {_nml}")
+        return task
+    if args.task == "labs_only":
+        return LabsOnlyMIMIC4(window_hours=args.observation_window_hours)
+    if args.task in ("cxr_only", "cxr_labs", "cxr_notes_labs"):
+        return CXRMultimodalMIMIC4(
+            window_hours=args.observation_window_hours,
+            include_labs=args.task in ("cxr_labs", "cxr_notes_labs"),
+            include_notes=args.task == "cxr_notes_labs",
+            frontal_only=not args.cxr_all_views,
+            image_size=args.cxr_image_size,
+            max_images=args.cxr_max_images,
+            note_source=args.note_source,
         )
     raise ValueError(f"Unknown task: {args.task}")
 
 
-def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any, str]:
+    """Split by patient, falling back to by-sample only if that yields nothing.
+
+    The fallback is leaky: a patient with several admissions can then land in
+    both train and test, which inflates the metrics. It only triggers on tiny
+    cohorts, but it must not trigger silently, so the mode is returned and
+    recorded alongside the run's results.
+    """
     train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
     if len(train_ds) == 0 or len(test_ds) == 0:
+        warnings.warn(
+            "split_by_patient produced an empty split, falling back to "
+            "split_by_sample. The same patient may now appear in train and "
+            "test, so these metrics are optimistic and not comparable to "
+            "patient-split runs.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
         train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
-    return train_ds, val_ds, test_ds
-
-
-def _build_model(args: argparse.Namespace, sample_dataset: Any):
+        return train_ds, val_ds, test_ds, "by_sample_fallback_leaky"
+    return train_ds, val_ds, test_ds, "by_patient"
+
+
+def _resolve_finetune_mode(args: argparse.Namespace) -> str:
+    """--freeze-encoder is a back-compat alias for --text-finetune-mode frozen."""
+    return "frozen" if args.freeze_encoder else args.text_finetune_mode
+
+
+def _build_model(
+    args: argparse.Namespace,
+    sample_dataset: Any,
+    numeric_standardizers: dict[str, Any] | None = None,
+):
+    finetune_mode = _resolve_finetune_mode(args)
+    field_embeddings = None
+    if "cxr" in sample_dataset.input_processors:
+        # Reuse the repository's VisionEmbeddingModel in the unified image
+        # branch; only the temporal alignment is supplied by unified.py.
+        field_embeddings = {
+            "cxr": VisionEmbeddingModel(
+                dataset=sample_dataset,
+                embedding_dim=args.embedding_dim,
+                patch_size=16,
+                backbone="patch",
+                pretrained=False,
+            )
+        }
     unified = UnifiedMultimodalEmbeddingModel(
         processors=sample_dataset.input_processors,
         embedding_dim=args.embedding_dim,
-        freeze_text_encoder=args.freeze_encoder,
+        field_embeddings=field_embeddings,
+        text_finetune_mode=finetune_mode,
+        normalize_content=not getattr(args, "no_normalize_content", False),
+        numeric_standardizers=numeric_standardizers,
+        cache_frozen_text=not getattr(args, "no_text_cache", False),
     )
 
     if args.model == "mlp":
@@ -228,6 +355,108 @@ def _write_predictions(
             )
 
 
+def _load_pretrained_weights(model, ckpt_path: str) -> None:
+    """Load SSL pretraining weights into a supervised model.
+
+    Delegates to ``BaseModel.load_pretrained_state_dict``, which performs the
+    architecture-specific key mapping (downstream models name the unified
+    backbone ``_unified_backbone`` / ``_unified_jamba`` / ``_unified_blocks``)
+    and REQUIRES full backbone coverage.
+
+    The previous implementation here built ``model.state_dict()``, overwrote
+    whichever checkpoint keys happened to map, and called ``strict=False``.
+    Because the untouched tensors were already present, PyTorch reported no
+    missing keys, so a jamba/mamba checkpoint that matched only ~6 of ~30
+    backbone tensors trained on a largely RANDOM backbone while looking
+    perfectly healthy. That silently corrupted real Table 2 cells.
+    """
+    print(f"[pretrain] Loading checkpoint from {ckpt_path}")
+    state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    try:
+        stats = model.load_pretrained_state_dict(state)
+    except ValueError as exc:
+        # Surface this as a run-level failure naming the checkpoint. A partial
+        # unified backbone must abort the job rather than train on random
+        # weights, and the operator needs to know WHICH checkpoint was bad.
+        raise RuntimeError(
+            f"Refusing to train on a partial unified backbone from {ckpt_path}: {exc}"
+        ) from exc
+    print(
+        "[pretrain] backbone {}/{} tensors, embedding {} matched; "
+        "uninitialised: {}".format(
+            stats["backbone_matched"], stats["backbone_target"],
+            stats["embedding_matched"], stats["missing_keys"][:6],
+        )
+    )
+
+def _note_availability_report(train_ds, sample_limit: int = 4000) -> dict:
+    """Measure how often a note is actually present, and whether that leaks.
+
+    A missing note is represented by a fixed placeholder embedding, so
+    "has a real note" is trivially learnable. Measured on MIMIC-IV, mortality was
+    5.67% where a note existed against 1.37% where it did not, a 4.1x gap: note
+    AVAILABILITY carries outcome signal with no clinical content behind it. That
+    confound has to be visible on every run rather than rediscovered, so it is
+    measured on the TRAIN split (never test) and recorded in run_config.json.
+    """
+    import numpy as np
+
+    total = len(train_ds)
+    if total == 0:
+        return {}
+    # Stride across the split rather than taking a prefix: samples are grouped by
+    # patient, so the first N are not representative of note availability.
+    step = max(1, total // sample_limit)
+    indices = range(0, total, step)
+    present, labels = [], []
+    for i in indices:
+        try:
+            sample = train_ds[i]
+        except Exception:
+            break
+        field = sample.get("admission_note_times")
+        if field is None:
+            return {}
+        try:
+            mask = field["mask"] if isinstance(field, dict) else field[1]
+            mask = torch.as_tensor(mask)
+            if mask.ndim == 1:
+                mask = mask.unsqueeze(0)
+            # content tokens = attention mask less [CLS] and [SEP]
+            content = int((mask.sum(dim=1) - 2).clamp(min=0).max())
+        except Exception:
+            return {}
+        present.append(content > 5)
+        labels.append(float(sample.get("mortality", 0)))
+    if not present or all(present) or not any(present):
+        return {"note_present_rate": float(np.mean(present)) if present else None,
+                "n_inspected": len(present)}
+    present = np.array(present); labels = np.array(labels)
+    report = {
+        "n_inspected": int(len(present)),
+        "note_present_rate": float(present.mean()),
+        "mortality_with_note": float(labels[present].mean()),
+        "mortality_without_note": float(labels[~present].mean()),
+    }
+    ratio = (report["mortality_with_note"] /
+             max(report["mortality_without_note"], 1e-9))
+    report["mortality_ratio_present_vs_absent"] = float(ratio)
+    print(
+        f"[note-availability] {100*report['note_present_rate']:.1f}% of train "
+        f"samples carry a real note; mortality {report['mortality_with_note']:.4f} "
+        f"with vs {report['mortality_without_note']:.4f} without ({ratio:.1f}x)."
+    )
+    if ratio > 1.5 or ratio < 0.67:
+        warnings.warn(
+            f"Note availability is {ratio:.1f}x associated with the outcome. The "
+            "missing-note placeholder is a constant embedding, so this is "
+            "learnable signal with no clinical content. Report it, restrict to "
+            "complete cases, or model missingness explicitly.",
+            RuntimeWarning, stacklevel=2,
+        )
+    return report
+
+
 def _compute_pos_weight(train_ds, label_key: str = "mortality") -> float:
     """Count pos/neg in train_ds and return n_neg/n_pos for BCE pos_weight."""
     n_pos = n_neg = 0
@@ -260,7 +489,56 @@ def run(args: argparse.Namespace) -> Path:
             "Task produced zero samples. Check roots/tables or adjust settings."
         )
 
-    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+    split_seed = getattr(args, "split_seed", None)
+    split_seed_pinned = split_seed is not None
+    if split_seed is None:
+        split_seed = args.seed
+        # Measured at full scale on labs_only: letting the split follow the seed
+        # gives sd(PR-AUC) 0.0236, versus 0.0042 with the split pinned, so test-set
+        # composition contributes ~31x the variance of initialisation. Cells
+        # compared across different splits are unpaired and far less sensitive.
+        warnings.warn(
+            "--split-seed not set, so the patient split follows --seed and every "
+            "seed draws a different test set. For an ablation, pin --split-seed "
+            "across compared cells: it is roughly a 5x sensitivity gain for no "
+            "extra compute.",
+            RuntimeWarning, stacklevel=2,
+        )
+    print(f"[split] patient split seed={split_seed} "
+          f"({'pinned' if split_seed_pinned else 'follows --seed'}); "
+          f"training seed={args.seed}")
+    train_ds, val_ds, test_ds, split_mode = _split_dataset(sample_dataset, seed=split_seed)
+    if getattr(args, "pretrained_ckpt", None) and split_mode != "by_patient":
+        raise ValueError(
+            "Pretrained comparisons require a non-empty patient-level split. "
+            "Refusing the sample-level fallback because the checkpoint's train "
+            "statistics could then include a downstream test patient."
+        )
+
+    # Fit before any outcome-dependent resampling and exclusively on the train
+    # subset.  ``SampleDataset.subset`` exposes only its selected indices here;
+    # LabStandardizer iterates this object, never ``sample_dataset``.
+    note_availability = _note_availability_report(train_ds)
+
+    numeric_standardizers: dict[str, Any] = {}
+    if "labs" in sample_dataset.input_processors and not getattr(
+        args, "no_lab_standardization", False
+    ):
+        if "labs_mask" not in sample_dataset.input_processors:
+            raise RuntimeError("Lab standardisation requires the labs_mask observation field.")
+        lab_standardizer = fit_lab_standardizer(
+            train_ds,
+            value_field="labs",
+            fit_scope=lab_standardizer_fit_scope(train_ds, value_field="labs"),
+        )
+        numeric_standardizers["labs"] = lab_standardizer
+        print(
+            "[lab-standardization] fitted on train split only: "
+            f"counts={lab_standardizer.observed_count.tolist()} "
+            f"mean={lab_standardizer.mean.tolist()} std={lab_standardizer.std.tolist()}"
+        )
+    elif "labs" in sample_dataset.input_processors:
+        print("[lab-standardization] disabled; reproducing raw-lab baseline.")
 
     label_key = list(sample_dataset.output_schema.keys())[0]
 
@@ -272,22 +550,41 @@ def run(args: argparse.Namespace) -> Path:
 
     if strategy == "undersample":
         ratio = args.balanced_ratio
-        print(f"[sampling] Undersampling negatives -> pos:neg 1:{ratio}")
+        print(f"[sampling] Undersampling negatives → pos:neg 1:{ratio}")
         train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key)
         print(f"[sampling] Training size after undersample: {len(train_ds)}")
 
     elif strategy == "oversample":
         ratio = args.balanced_ratio
-        print(f"[sampling] Oversampling positives -> pos:neg 1:{ratio}")
+        print(f"[sampling] Oversampling positives → pos:neg 1:{ratio}")
         train_ds = sample_oversample(train_ds, ratio=ratio, seed=args.seed, label_key=label_key)
         print(f"[sampling] Training size after oversample: {len(train_ds)}")
 
     elif strategy == "weighted":
-        print("[sampling] Weighted resampling (class-proportional, with replacement)")
+        print("[sampling] Weighted resampling (class-proportional, with replacement, no external sampler)")
         train_ds = sample_weighted(train_ds, seed=args.seed, label_key=label_key)
         print(f"[sampling] Training size after weighted resample: {len(train_ds)}")
 
-    model = _build_model(args, sample_dataset)
+    model = _build_model(args, sample_dataset, numeric_standardizers)
+
+    # Load pretrained SSL weights if requested.
+    if getattr(args, "pretrained_ckpt", None):
+        _load_pretrained_weights(model, args.pretrained_ckpt)
+    if args.numeric_input_stats_path:
+        model.embedding_model.capture_numeric_encoder_input_stats(
+            args.numeric_input_stats_path, field_name="labs"
+        )
+        # Use the same deterministic train-split batch as the pretraining
+        # audit. This records the true projection input while avoiding a
+        # misleading difference caused only by independent shuffling.
+        audit_batch = next(iter(get_dataloader(
+            train_ds, batch_size=args.batch_size, shuffle=False,
+            num_workers=args.loader_num_workers,
+        )))
+        with torch.no_grad():
+            model(**audit_batch)
+        if not Path(args.numeric_input_stats_path).is_file():
+            raise RuntimeError("Numeric-input audit hook did not produce its artifact.")
 
     # Apply class-imbalance correction via BCE pos_weight.
     # pos_weight = n_neg / n_pos so the rare positive class gets proportionally
@@ -300,18 +597,64 @@ def run(args: argparse.Namespace) -> Path:
     print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.")
     model._pos_weight = torch.tensor([pw_value], dtype=torch.float32)
 
-    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    loader_kwargs = {
+        "num_workers": args.loader_num_workers,
+        "pin_memory": args.pin_memory,
+        "persistent_workers": (
+            args.persistent_workers and args.loader_num_workers > 0
+        ),
+        "prefetch_factor": (
+            args.prefetch_factor if args.loader_num_workers > 0 else None
+        ),
+    }
+    train_loader = get_dataloader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        **loader_kwargs,
+    )
     val_loader = (
-        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        get_dataloader(
+            val_ds,
+            batch_size=args.batch_size,
+            shuffle=False,
+            **loader_kwargs,
+        )
         if len(val_ds) > 0
         else None
     )
     test_loader = (
-        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        get_dataloader(
+            test_ds,
+            batch_size=args.batch_size,
+            shuffle=False,
+            **loader_kwargs,
+        )
         if len(test_ds) > 0
         else None
     )
 
+    # Which split the reported predictions come from. Falling back to val, or
+    # worse to train, silently reports held-in performance as if it were test,
+    # so resolve it once here and record it alongside the results.
+    if test_loader is not None:
+        inference_loader, eval_split = test_loader, "test"
+    elif val_loader is not None:
+        inference_loader, eval_split = val_loader, "val"
+        warnings.warn(
+            "No test split available; reporting predictions from the VALIDATION "
+            "split. These are not test metrics.",
+            RuntimeWarning, stacklevel=2,
+        )
+    else:
+        inference_loader, eval_split = train_loader, "train"
+        warnings.warn(
+            "No test or validation split available; reporting predictions from "
+            "the TRAINING split. These metrics are held-in and meaningless as a "
+            "generalisation estimate.",
+            RuntimeWarning, stacklevel=2,
+        )
+
     # Experiment name encodes model + seed for easy log separation
     exp_name = f"{args.model}_seed{args.seed}"
     output_dir = Path(args.output_dir)
@@ -323,6 +666,8 @@ def run(args: argparse.Namespace) -> Path:
         enable_logging=True,
         output_path=str(output_dir),
         exp_name=exp_name,
+        use_amp=args.use_amp,
+        amp_dtype=args.amp_dtype,
     )
 
     # BottleneckTransformer is more fragile on full MIMIC-IV with no warmup.
@@ -352,7 +697,37 @@ def run(args: argparse.Namespace) -> Path:
 
     optimizer_params["lr"] = effective_lr
 
+    # Record the resolved conditions, not the raw flags: text_finetune_mode and
+    # the learning rate are both derived, so the CLI alone does not identify the
+    # run. Without this the artifacts cannot say whether the encoder was frozen.
+    write_run_config(
+        str(output_dir / exp_name),
+        {
+            **vars(args),
+            "resolved_text_finetune_mode": _resolve_finetune_mode(args),
+            "resolved_lr": effective_lr,
+            "resolved_max_grad_norm": effective_max_grad_norm,
+            "split_mode": split_mode,
+            "eval_split": eval_split,
+            "split_seed_pinned": split_seed_pinned,
+            "note_availability": note_availability,
+            "n_train": len(train_ds),
+            "n_val": len(val_ds),
+            "n_test": len(test_ds),
+        },
+    )
+
     if args.epochs > 0 and len(train_ds) > 0:
+        # PR-AUC/ROC-AUC are undefined for a single-class validation fold.
+        # Full MIMIC patient splits contain both labels; tiny real-data demo
+        # runs do not, so use finite validation loss for checkpoint selection.
+        val_labels = {
+            int(float(val_ds[i][label_key])) for i in range(len(val_ds))
+        }
+        monitor = "pr_auc" if len(val_labels) == 2 else "loss"
+        monitor_criterion = "max" if monitor == "pr_auc" else "min"
+        if monitor != "pr_auc":
+            print("[monitor] Validation fold has one class; selecting checkpoints by loss.")
         trainer.train(
             train_dataloader=train_loader,
             val_dataloader=val_loader,
@@ -360,12 +735,13 @@ def run(args: argparse.Namespace) -> Path:
             optimizer_params=optimizer_params,
             weight_decay=args.weight_decay,
             max_grad_norm=effective_max_grad_norm,
-            monitor="pr_auc",
+            monitor=monitor,
+            monitor_criterion=monitor_criterion,
             load_best_model_at_last=True,
             patience=args.patience,
+            encoder_lr=args.encoder_lr,
         )
 
-    inference_loader = test_loader or val_loader or train_loader
     y_true, y_prob, _, patient_ids = trainer.inference(
         inference_loader, return_patient_ids=True
     )
@@ -381,13 +757,28 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--ehr-root", type=str, required=True)
     parser.add_argument("--note-root", type=str, default=None)
+    parser.add_argument("--cxr-root", type=str, default=None)
     parser.add_argument("--cache-dir", type=str, default=None)
     parser.add_argument("--output-dir", type=str, default="./output/unified_e2e")
+    parser.add_argument(
+        "--numeric-input-stats-path", type=str, default=None,
+        help="Optional JSON artifact: first lab tensor entering the numeric encoder.",
+    )
+    parser.add_argument(
+        "--pretrained-ckpt",
+        type=str,
+        default=None,
+        help=(
+            "Path to a SSL pretraining checkpoint (e.g., from "
+            "scripts/pretrain_ssl.py).  Loads embedding_model and backbone "
+            "weights into the downstream model."
+        ),
+    )
 
     parser.add_argument(
         "--task",
         type=str,
-        choices=["icd_labs", "clinical_notes_icd_labs"],
+        choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "notes_only", "labs_only", "cxr_only", "cxr_labs", "cxr_notes_labs"],
         default="stagenet",
         help=(
             "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. "
@@ -430,7 +821,32 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--weight-decay", type=float, default=0.0)
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument("--num-workers", type=int, default=1)
+    parser.add_argument(
+        "--loader-num-workers",
+        type=int,
+        default=0,
+        help="Worker processes for runtime batch loading/collation.",
+    )
+    parser.add_argument("--pin-memory", action="store_true", default=False)
+    parser.add_argument("--persistent-workers", action="store_true", default=False)
+    parser.add_argument("--prefetch-factor", type=int, default=2)
+    amp_group = parser.add_mutually_exclusive_group()
+    amp_group.add_argument("--use-amp", dest="use_amp", action="store_true")
+    amp_group.add_argument("--no-amp", dest="use_amp", action="store_false")
+    parser.set_defaults(use_amp=False)
+    parser.add_argument(
+        "--amp-dtype",
+        choices=["bf16", "fp16"],
+        default="bf16",
+        help="AMP compute dtype; bf16 is recommended on A100/H100.",
+    )
     parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--split-seed",
+        type=int,
+        default=None,
+        help="Patient split seed. Defaults to --seed for backward compatibility.",
+    )
     parser.add_argument("--patience", type=int, default=None)
     parser.add_argument(
         "--dev",
@@ -457,6 +873,12 @@ def parse_args() -> argparse.Namespace:
 
     # Task-specific
     parser.add_argument("--observation-window-hours", type=int, default=24)
+    parser.add_argument(
+        "--cxr-all-views", action="store_true", default=False,
+        help="Keep non-frontal CXR views; default restricts to PA/AP frontal images.",
+    )
+    parser.add_argument("--cxr-image-size", type=int, default=224)
+    parser.add_argument("--cxr-max-images", type=int, default=4)
     parser.add_argument(
         "--icd-codes",
         action="store_true",
@@ -468,6 +890,22 @@ def parse_args() -> argparse.Namespace:
             "Enable only for ablation / legacy comparison experiments."
         ),
     )
+    parser.add_argument(
+        "--discharge-note-policy",
+        choices=["extraction", "charttime"],
+        default="extraction",
+        help="How the observation window applies to discharge summaries. "
+             "extraction (default, Lee et al. 2023): retrieve across the "
+             "admission and let admission-context section extraction be the "
+             "temporal control. charttime: strictly non-anticipative, but "
+             "drops the summary for ~90%% of admissions.",
+    )
+    parser.add_argument(
+        "--no-text-cache",
+        action="store_true",
+        help="Disable the frozen-text [CLS] cache. Diagnostic: isolates the cache "
+             "as a cause when a run fails to optimise.",
+    )
     parser.add_argument(
         "--freeze-encoder",
         action="store_true",
@@ -475,8 +913,44 @@ def parse_args() -> argparse.Namespace:
         help=(
             "Freeze pretrained BERT text encoder weights and train only the "
             "downstream backbone (MLP/RNN/Transformer head + projection layer). "
-            "Reduces VRAM by ~50% for the text branch; useful when GPU memory "
-            "is limited or for faster iteration on backbone architectures."
+            "Reduces VRAM by ~50%% for the text branch; useful when GPU memory "
+            "is limited or for faster iteration on backbone architectures. "
+            "Back-compat alias for --text-finetune-mode frozen."
+        ),
+    )
+    parser.add_argument(
+        "--text-finetune-mode",
+        type=str,
+        default="full",
+        help=(
+            "Text encoder fine-tuning regime: 'full' (train all encoder params), "
+            "'frozen' (train only the head/projection), 'topk:N' (unfreeze the top "
+            "N transformer layers, embeddings stay frozen), or 'lora:r' (rank-r "
+            "LoRA adapters on attention, base frozen; needs peft). "
+            "Overridden by --freeze-encoder when that flag is set."
+        ),
+    )
+    parser.add_argument(
+        "--encoder-lr",
+        type=float,
+        default=None,
+        help=(
+            "Discriminative learning rate for the pretrained text encoder. When "
+            "set, the encoder trains at this LR while the projection + downstream "
+            "head keep the base --lr. Recommended ~2e-5 for full/topk fine-tuning "
+            "to avoid destabilizing the encoder. Default None = single global LR."
+        ),
+    )
+    parser.add_argument(
+        "--note-source",
+        type=str,
+        default="discharge",
+        choices=["discharge", "radiology"],
+        help=(
+            "Which MIMIC note table to use for notes_labs. 'discharge' (default) "
+            "uses admission-context discharge sections; 'radiology' uses radiology "
+            "report Impression/Findings, concatenated per admission (mirrors the "
+            "multimodal-EHR benchmark). Pair with a RadBERT --tokenizer-model."
         ),
     )
     parser.add_argument(
@@ -508,6 +982,83 @@ def parse_args() -> argparse.Namespace:
             "Default: 1.0 (equal pos/neg). Used with undersample and oversample strategies."
         ),
     )
+    parser.add_argument(
+        "--tokenizer-model",
+        type=str,
+        default=None,
+        help=(
+            "Override the tokenizer/encoder for notes. Must be a HuggingFace model ID. "
+            "Default: None (uses task class default, emilyalsentzer/Bio_ClinicalBERT). "
+            "Changes the task cache UUID so different tokenizers use isolated caches. "
+            "Examples: microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext, "
+            "medicalai/ClinicalBERT, yikuan8/Clinical-Longformer."
+        ),
+    )
+    parser.add_argument(
+        "--no-normalize-content",
+        action="store_true",
+        default=False,
+        help=(
+            "Disable content normalization in the unified embedding, reproducing "
+            "pre-repair behaviour where text events were ~94%% patient-independent "
+            "constant while raw labs dominated. Use only to reproduce old runs."
+        ),
+    )
+    parser.add_argument(
+        "--no-lab-standardization",
+        action="store_true",
+        default=False,
+        help=(
+            "Disable train-split-only per-analyte lab z-scoring. Use only for "
+            "the raw-lab ablation; default standardises immediately before the "
+            "numeric projection and checkpoints the fitted statistics."
+        ),
+    )
+    parser.add_argument(
+        "--note-max-length",
+        type=int,
+        default=None,
+        help=(
+            "Note tokenization budget (TupleTimeTextProcessor.max_length; default "
+            "128, which truncates ~96%% of extracted discharge notes). BERT-family "
+            "encoders cap at 512 position embeddings; for longer budgets pass "
+            "--tokenizer-model yikuan8/Clinical-Longformer (4096) and expect to "
+            "need --batch-size 1."
+        ),
+    )
+    parser.add_argument(
+        "--text-normalize",
+        type=str,
+        default="none",
+        choices=["none", "punct", "stopwords", "both"],
+        help=(
+            "Cut tokens per note before tokenization. 'punct': strip punctuation "
+            "(decimal points and thousands separators inside numbers are kept, so "
+            "lab values survive). 'stopwords': drop common English stopwords. "
+            "'both': both. Changes task_name so each setting gets its own cache."
+        ),
+    )
+    parser.add_argument(
+        "--note-extraction",
+        type=str,
+        default="regex",
+        choices=[
+            "regex", "regex_priority", "compact", "tfidf",
+            "section_hpi", "section_cc", "section_pmh", "section_meds",
+            "section_social", "section_family", "section_allergies", "section_ros",
+            "lab_retrieval",
+        ],
+        help=(
+            "Note text extraction strategy. "
+            "'regex' (default): section headers, document order. "
+            "'regex_priority': HPI-first order (fixes 48%% truncation rate). "
+            "'compact': HPI + CC only — always fits in 512 tokens. "
+            "'tfidf': TF overlap paragraph retrieval, no headers required. "
+            "'section_<name>': single-section ablation (hpi/cc/pmh/meds/"
+            "social/family/allergies/ros). "
+            "Non-regex values isolate the task cache UUID."
+        ),
+    )
     parser.add_argument(
         "--sampling-strategy",
         type=str,
@@ -518,7 +1069,7 @@ def parse_args() -> argparse.Namespace:
             "'none': no resampling (default). "
             "'undersample': drop majority-class (neg) samples via sample_balanced(). "
             "'oversample': duplicate minority-class (pos) samples via sample_oversample(). "
-            "'weighted': class-proportional resampling w/ replacement via sample_weighted(). "
+            "'weighted': WeightedRandomSampler for batch-level balance without dataset modification. "
             "--balanced-sampling is a legacy alias for 'undersample'."
         ),
     )
@@ -528,7 +1079,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--rnn-layers", type=int, default=1)
     parser.add_argument("--bidirectional", action="store_true")
 
-    # Transformer / BottleneckTransformer shared
+    # Transformer / BottleneckTransformer shared.
+    # Standardized compute: 64-dim, 1 layer, 2 heads (head_dim 32) across all archs.
     parser.add_argument("--heads", type=int, default=4)
     parser.add_argument("--num-layers", type=int, default=2)
 
@@ -553,9 +1105,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--mamba-conv-kernel", type=int, default=4,
                         help="Causal conv kernel size for EHRMamba and JambaEHR blocks.")
     parser.add_argument("--jamba-transformer-layers", type=int, default=2,
-                        help="Number of Transformer (attention) layers in JambaEHR.")
+                        help="Number of Transformer (attention) layers in JambaEHR. "
+                             "Standard: 1 (a single Jamba block = 1 attn + 1 mamba).")
     parser.add_argument("--jamba-mamba-layers", type=int, default=6,
-                        help="Number of Mamba (SSM) layers in JambaEHR.")
+                        help="Number of Mamba (SSM) layers in JambaEHR. "
+                             "Standard: 1 (a single Jamba block = 1 attn + 1 mamba).")
 
     return parser.parse_args()
 
diff --git a/pyhealth/datasets/configs/mimic4_cxr.yaml b/pyhealth/datasets/configs/mimic4_cxr.yaml
index fb5045eef..bca6eb95e 100644
--- a/pyhealth/datasets/configs/mimic4_cxr.yaml
+++ b/pyhealth/datasets/configs/mimic4_cxr.yaml
@@ -5,7 +5,7 @@ tables:
     patient_id: "subject_id"
     timestamp:
       - "studydate"
-      - "studytime"
+      - "studytime_normalized"
     timestamp_format: "%Y%m%d%H%M%S"
     attributes:
       - "image_path"
diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py
index 089a6d2f3..9be17e1b8 100644
--- a/pyhealth/datasets/mimic4.py
+++ b/pyhealth/datasets/mimic4.py
@@ -191,6 +191,15 @@ def __init__(
         log_memory_usage(f"After initializing {dataset_name}")
 
     def prepare_metadata(self, root: str) -> None:
+        prepared_path = os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv")
+        # The prepared file holds absolute JPEG paths. To rewrite its 100+ MB
+        # CSV for every CXR experiment is needless shared-filesystem IO, and it
+        # races a concurrent reader, so reuse the file when it is present. An
+        # installation with raw metadata only still uses the builder below.
+        if os.path.exists(prepared_path):
+            header = pd.read_csv(prepared_path, nrows=1)
+            if "image_path" in header.columns:
+                return
         metadata = pd.read_csv(
             os.path.join(root, "mimic-cxr-2.0.0-metadata.csv.gz"), dtype=str
         )
@@ -217,9 +226,7 @@ def process_image_path(x):
 
         metadata["image_path"] = metadata.apply(process_image_path, axis=1)
 
-        metadata.to_csv(
-            os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv"), index=False
-        )
+        metadata.to_csv(prepared_path, index=False)
         return
 
 
diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py
index cadb479ce..d892252ca 100644
--- a/pyhealth/tasks/__init__.py
+++ b/pyhealth/tasks/__init__.py
@@ -74,5 +74,6 @@
     ClinicalNotesMIMIC4,
     ClinicalNotesICDLabsMIMIC4,
     ClinicalNotesICDLabsCXRMIMIC4,
+    CXRMultimodalMIMIC4,
 )
 from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task
diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py
index 21147303c..d7c944b5a 100644
--- a/pyhealth/tasks/multimodal_mimic4.py
+++ b/pyhealth/tasks/multimodal_mimic4.py
@@ -1447,3 +1447,200 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             record["icd_codes"] = (all_icd_times, all_icd_codes)
 
         return [record]
+
+
+class CXRMultimodalMIMIC4(BaseMultimodalMIMIC4Task):
+    """Admission-level MIMIC-IV mortality task with timestamped MIMIC-CXR.
+
+    MIMIC-CXR has no ``hadm_id``.  This task therefore links an image event to
+    an admission only when the *same subject's* ``StudyDate + StudyTime`` falls
+    inside that admission's observation interval.  This is intentionally done
+    in the task (instead of a loose subject-level merge) so studies acquired
+    after prediction time cannot enter the sample.
+
+    MIMIC-CXR and MIMIC-IV timestamps are both de-identified with the same
+    patient-specific date shift.  Consequently, ``metadata.timestamp -
+    admission.timestamp`` is a valid hours-from-admission time just like the
+    lab offsets used by :meth:`_collect_labs`.
+
+    The default is PA/AP frontal images only.  Frontal views are the standard
+    clinical CXR representation and exclude lateral/oblique acquisitions whose
+    geometry makes a single lightweight patch encoder substantially less
+    comparable.  Set ``frontal_only=False`` to retain every view.
+
+    One sample is emitted per admission with at least one eligible CXR.  This
+    preserves the actual CXR-to-admission linkage rather than injecting an
+    invalid blank image placeholder for unlinked admissions.  Patient-level
+    splitting remains required downstream to avoid patient overlap.
+
+    Args:
+        window_hours: Observation window from admission. Defaults to 24 h.
+        include_labs: Include the existing timestamped 10-dimensional lab
+            vectors in the same admission window.
+        include_notes: Include existing admission-context notes in the same
+            window. Requires the corresponding MIMIC-IV note table.
+        frontal_only: Keep PA/AP images only. Defaults to ``True``.
+        image_size: Square loader resize. Defaults to 224.
+        max_images: Cap images per admission after chronological sorting. The
+            most recent images are retained, matching ``TimeImageProcessor``.
+        note_source: ``"discharge"`` or ``"radiology"`` when notes are used.
+    """
+
+    task_name: str = "CXRMultimodalMIMIC4"
+    output_schema: Dict[str, str] = {"mortality": "binary"}
+
+    _CXR_SCHEMA: ClassVar[Dict[str, Union[str, Tuple[str, Dict]]]] = {
+        # MIMIC-CXR JPEGs are grayscale.  Decoding to one channel avoids a
+        # gratuitous 3x input expansion while VisionEmbeddingModel/PatchEmbedding
+        # automatically receives the processor's inferred channel count.
+        "cxr": ("time_image", {"image_size": 224, "mode": "L", "max_images": 4}),
+    }
+    _NOTES_SCHEMA: ClassVar[Tuple[str, Dict]] = (
+        "tuple_time_text",
+        {
+            "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT",
+            "max_length": 512,
+            "type_tag": "note",
+        },
+    )
+
+    def __init__(
+        self,
+        window_hours: Optional[float] = 24,
+        include_labs: bool = False,
+        include_notes: bool = False,
+        frontal_only: bool = True,
+        image_size: int = 224,
+        max_images: Optional[int] = 4,
+        note_source: str = "discharge",
+    ) -> None:
+        if window_hours is not None and window_hours < 0:
+            raise ValueError("window_hours must be non-negative or None.")
+        if image_size <= 0:
+            raise ValueError("image_size must be positive.")
+        if max_images is not None and max_images <= 0:
+            raise ValueError("max_images must be positive or None.")
+        if note_source not in {"discharge", "radiology"}:
+            raise ValueError(f"Unsupported note_source: {note_source!r}")
+
+        super().__init__(window_hours=window_hours)
+        # Included in the task cache UUID; increment on emitted-record changes.
+        self.cxr_pipeline_cache_version = 1
+        self.include_labs = include_labs
+        self.include_notes = include_notes
+        self.frontal_only = frontal_only
+        self.image_size = image_size
+        self.max_images = max_images
+        self.note_source = note_source
+
+        schema = dict(self._CXR_SCHEMA)
+        schema["cxr"] = (
+            "time_image",
+            {"image_size": image_size, "mode": "L", "max_images": max_images},
+        )
+        if include_labs:
+            schema["labs"] = ("stagenet_tensor", {})
+            schema["labs_mask"] = ("stagenet_tensor", {})
+        if include_notes:
+            schema["admission_note_times"] = self._NOTES_SCHEMA
+        self.input_schema = schema
+
+        suffix = ["frontal" if frontal_only else "allviews"]
+        if include_labs:
+            suffix.append("labs")
+        if include_notes:
+            suffix.append(f"notes_{note_source}")
+        self.task_name = "CXRMultimodalMIMIC4_" + "_".join(suffix)
+
+    def _cxr_window_end(self, admission_time: datetime, admission: Any) -> datetime:
+        """Return this admission's leakage-safe CXR/lab/note cutoff."""
+        dischtime = self._parse_datetime(getattr(admission, "dischtime", None))
+        if dischtime is None or dischtime < admission_time:
+            dischtime = admission_time
+        if self.window_hours is None:
+            return dischtime
+        return min(dischtime, admission_time + timedelta(hours=self.window_hours))
+
+    def _collect_admission_cxr(
+        self,
+        patient: Any,
+        admission_time: datetime,
+        end_time: datetime,
+    ) -> Tuple[List[str], List[float]]:
+        """Collect valid CXR paths and real hours-from-admission timestamps."""
+        events = patient.get_events(
+            event_type="metadata", start=admission_time, end=end_time
+        )
+        paths: List[str] = []
+        times: List[float] = []
+        seen_dicom: set[str] = set()
+        for event in events:
+            try:
+                view = str(getattr(event, "viewposition", "")).upper().strip()
+                if self.frontal_only and view not in {"PA", "AP"}:
+                    continue
+                path = str(event.image_path)
+                dicom_id = str(getattr(event, "dicom_id", path))
+                if not path or dicom_id in seen_dicom:
+                    continue
+                seen_dicom.add(dicom_id)
+                paths.append(path)
+                times.append(
+                    self._to_hours((event.timestamp - admission_time).total_seconds())
+                )
+            except AttributeError:
+                # A malformed metadata record cannot become an image event.
+                continue
+        return paths, times
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:
+        if not patient.get_events(event_type="patients"):
+            return []
+        records: List[Dict[str, Any]] = []
+        for admission in patient.get_events(event_type="admissions"):
+            admission_time = admission.timestamp
+            # Discharge-note retrieval spans the admission (the section
+            # extraction is the temporal control), so this loop needs dischtime.
+            admission_dischtime = (
+                self._parse_datetime(getattr(admission, "dischtime", None))
+                or admission_time
+            )
+            end_time = self._cxr_window_end(admission_time, admission)
+            cxr_paths, cxr_times = self._collect_admission_cxr(
+                patient, admission_time, end_time
+            )
+            # Do not fabricate an image sentinel: an unreadable blank path would
+            # fail the image processor and, more importantly, hide cohort shift.
+            if not cxr_paths:
+                continue
+
+            record: Dict[str, Any] = {
+                "patient_id": patient.patient_id,
+                "hadm_id": str(getattr(admission, "hadm_id", "")),
+                "cxr": (cxr_paths, cxr_times),
+                "mortality": int(getattr(admission, "hospital_expire_flag", 0) in (1, "1")),
+                "window_start": admission_time,
+                "window_end": end_time,
+            }
+            if self.include_labs:
+                lab_times, lab_values, lab_masks = self._collect_labs(
+                    patient=patient, admission_time=admission_time, end_time=end_time
+                )
+                record["labs"] = (lab_times, lab_values)
+                record["labs_mask"] = (lab_times, lab_masks)
+            if self.include_notes:
+                # Use the admission-context section extraction that main
+                # already has. A discharge summary is written with knowledge of
+                # the outcome, so only the admission-context sections enter a
+                # sample.
+                texts, note_times = self._collect_notes(
+                    patient,
+                    self.note_source,
+                    getattr(admission, "hadm_id", None),
+                    admission_time,
+                    end_time=end_time,
+                    section_headers=self.DISCHARGE_CLINICAL_HEADERS,
+                )
+                record["admission_note_times"] = (texts, note_times)
+            records.append(record)
+        return records
diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py
index fc264a3af..a774e0c81 100644
--- a/pyhealth/trainer.py
+++ b/pyhealth/trainer.py
@@ -61,6 +61,30 @@ def get_metrics_fn(mode: str) -> Callable:
         raise ValueError(f"Mode {mode} is not supported")
 
 
+def is_text_pathway_param(name: str, text_fields, frozen_text_fields=()) -> bool:
+    """Whether a parameter belongs to the text pathway, for discriminative LR.
+
+    ``encoder_lr`` exists to give a PRETRAINED encoder a gentler learning rate
+    than the randomly initialised layers around it, so a projection normally
+    keeps the base rate.
+
+    The exception is a frozen encoder. Then every ``encoders.*`` parameter has
+    ``requires_grad=False``, the group is empty, and ``encoder_lr`` controls
+    nothing, while the only trainable text parameters, ``projections.*``, keep
+    the base rate. For those fields the projection IS the text pathway, so it
+    joins the group.
+    """
+    frozen = set(frozen_text_fields or ())
+    return any(
+        name.startswith(f"embedding_model.encoders.{field}.")
+        or (
+            field in frozen
+            and name.startswith(f"embedding_model.projections.{field}.")
+        )
+        for field in text_fields
+    )
+
+
 class Trainer:
     """Trainer for PyTorch models.
 
@@ -140,6 +164,7 @@ def train(
         accumulation_steps: int = 1,
         use_amp: bool = False,
         amp_dtype: str = "bf16",
+        encoder_lr: Optional[float] = None,
     ):
         """Trains the model.
 
@@ -196,16 +221,49 @@ def train(
         # set optimizer
         param = list(self.model.named_parameters())
         no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
+        def _decayed(n):
+            return not any(nd in n for nd in no_decay)
+
         optimizer_grouped_parameters = [
             {
-                "params": [p for n, p in param if not any(nd in n for nd in no_decay)],
+                "params": [p for n, p in param if _decayed(n)],
                 "weight_decay": weight_decay,
             },
             {
-                "params": [p for n, p in param if any(nd in n for nd in no_decay)],
+                "params": [p for n, p in param if not _decayed(n)],
                 "weight_decay": 0.0,
             },
         ]
+        if encoder_lr is not None:
+            # A pretrained text encoder needs a gentler rate than the randomly
+            # initialised layers around it. Everything else keeps the base rate
+            # from optimizer_params.
+            embedding_model = getattr(self.model, "embedding_model", None)
+            modality_types = getattr(embedding_model, "modality_types", {})
+            text_fields = {
+                field
+                for field, modality in modality_types.items()
+                if getattr(modality, "value", modality) == "text"
+            }
+            if not text_fields:
+                raise ValueError(
+                    "encoder_lr was set, but the model has no text encoder"
+                )
+            frozen_text_fields = getattr(embedding_model, "_frozen_text_fields", set())
+
+            def _is_encoder(n):
+                return is_text_pathway_param(n, text_fields, frozen_text_fields)
+
+            optimizer_grouped_parameters = [
+                {"params": [p for n, p in param if _is_encoder(n) and _decayed(n)],
+                 "weight_decay": weight_decay, "lr": encoder_lr},
+                {"params": [p for n, p in param if _is_encoder(n) and not _decayed(n)],
+                 "weight_decay": 0.0, "lr": encoder_lr},
+                {"params": [p for n, p in param if not _is_encoder(n) and _decayed(n)],
+                 "weight_decay": weight_decay},
+                {"params": [p for n, p in param if not _is_encoder(n) and not _decayed(n)],
+                 "weight_decay": 0.0},
+            ]
         optimizer = optimizer_class(optimizer_grouped_parameters, **optimizer_params)
 
         # initialize
@@ -300,10 +358,18 @@ def train(
             if self.exp_path is not None:
                 self.save_ckpt(os.path.join(self.exp_path, "last.ckpt"))
 
+            # An epoch mean cannot show the difference between a run that
+            # starts badly and a run that becomes worse inside the epoch.
+            # Record the trajectory so a divergence is visible without a rerun.
+            _head = training_loss[: min(100, len(training_loss))]
+            _tail = training_loss[-min(100, len(training_loss)):]
             epoch_record: Dict = {
                 "epoch": epoch,
                 "global_step": global_step,
                 "train_loss": sum(training_loss) / len(training_loss),
+                "train_loss_first_step": round(training_loss[0], 6),
+                "train_loss_first100": round(sum(_head) / len(_head), 6),
+                "train_loss_last100": round(sum(_tail) / len(_tail), 6),
                 "epoch_time_s": round(epoch_time, 3),
                 **{f"train_{k}": v for k, v in vram.items()},
             }
diff --git a/pyhealth/utils.py b/pyhealth/utils.py
index efbee3977..f6cd6ff4c 100644
--- a/pyhealth/utils.py
+++ b/pyhealth/utils.py
@@ -1,7 +1,9 @@
+import hashlib
 import json
 import os
 import pickle
 import random
+import subprocess
 import contextlib
 
 import numpy as np
@@ -43,6 +45,82 @@ def save_json(data, filename):
     with open(filename, "w") as f:
         json.dump(data, f)
 
+
+def _git_revision():
+    repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    try:
+        commit = subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], cwd=repo, stderr=subprocess.DEVNULL
+        ).decode().strip()
+        dirty = bool(subprocess.check_output(
+            ["git", "status", "--porcelain"], cwd=repo, stderr=subprocess.DEVNULL
+        ).decode().strip())
+        return {"commit": commit, "dirty": dirty}
+    except Exception:
+        return {"commit": None, "dirty": None}
+
+
+def _source_digest():
+    """Hash the package source so code identity survives a non-git deploy.
+
+    Cluster runs typically execute from an unpacked tarball rather than a
+    clone, so the git lookup returns nothing exactly where provenance matters
+    most. Hashing the sources keeps "which code produced this result"
+    answerable either way.
+    """
+    package = os.path.dirname(os.path.abspath(__file__))
+    digest = hashlib.sha256()
+    try:
+        for root, dirs, files in os.walk(package):
+            dirs[:] = sorted(d for d in dirs if d != "__pycache__")
+            for name in sorted(files):
+                if not name.endswith(".py"):
+                    continue
+                path = os.path.join(root, name)
+                digest.update(os.path.relpath(path, package).encode())
+                with open(path, "rb") as f:
+                    digest.update(f.read())
+        return digest.hexdigest()
+    except OSError:
+        return None
+
+
+def _jsonable(value):
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    if isinstance(value, (list, tuple)):
+        return [_jsonable(v) for v in value]
+    if isinstance(value, dict):
+        return {str(k): _jsonable(v) for k, v in value.items()}
+    return str(value)
+
+
+def write_run_config(exp_path, config):
+    """Persist the resolved run configuration next to the run's metrics.
+
+    ``metrics_history.json`` records what a run scored but not the conditions
+    that produced it, so a frozen-encoder run and a fine-tuned one are
+    indistinguishable once the job's stdout is gone. Record the resolved
+    settings, not the raw flags, so derived conditions are recoverable.
+    """
+    record = {
+        "config": {str(k): _jsonable(v) for k, v in config.items()},
+        "git": _git_revision(),
+        "source_sha256": _source_digest(),
+        "torch": torch.__version__,
+    }
+    os.makedirs(exp_path, exist_ok=True)
+    path = os.path.join(exp_path, "run_config.json")
+    tmp = f"{path}.tmp.{os.getpid()}"
+    try:
+        with open(tmp, "w") as f:
+            json.dump(record, f, indent=2, sort_keys=True)
+        os.replace(tmp, path)
+    finally:
+        if os.path.exists(tmp):
+            os.remove(tmp)
+    return path
+
 @contextlib.contextmanager
 def set_env(**environ):
     """
diff --git a/tests/test_run_provenance_and_pathways.py b/tests/test_run_provenance_and_pathways.py
new file mode 100644
index 000000000..6d686a818
--- /dev/null
+++ b/tests/test_run_provenance_and_pathways.py
@@ -0,0 +1,215 @@
+"""Regression tests for run provenance and the two silent fallbacks.
+
+Each defect below made a measurement mean something other than what it said. A
+run reported validation performance as test performance. A patient split became
+a sample split, which leaks. A frozen-encoder run and a fine-tuned run left
+identical artefacts on disk, so the condition of an earlier result could not be
+recovered.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Run provenance
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def test_the_run_configuration_is_written_beside_the_metrics(tmp_path):
+    from pyhealth.utils import write_run_config
+
+    path = write_run_config(str(tmp_path), {"task": "notes_labs", "seed": 42})
+    record = json.loads(Path(path).read_text())
+
+    assert Path(path).name == "run_config.json"
+    assert record["config"]["task"] == "notes_labs"
+    assert record["config"]["seed"] == 42
+
+
+def test_code_identity_survives_a_run_from_an_unpacked_archive(tmp_path):
+    """A cluster run starts from a tarball, where `git rev-parse` gives nothing.
+    A digest of the package source identifies the code in that case.
+    """
+    from pyhealth.utils import write_run_config
+
+    record = json.loads(Path(write_run_config(str(tmp_path), {})).read_text())
+
+    assert record["source_sha256"], "no code identity was recorded"
+    assert len(record["source_sha256"]) >= 16
+    assert "torch" in record
+
+
+def test_the_digest_is_stable_for_unchanged_source():
+    """Two runs of the same code must record the same digest, otherwise the
+    digest cannot show that a set of arms shared one version of the code.
+    """
+    from pyhealth import utils
+
+    assert utils._source_digest() == utils._source_digest()
+
+
+def test_a_value_that_json_cannot_hold_does_not_lose_the_whole_record(tmp_path):
+    """A configuration holds Paths, enums and devices. If one of them raises,
+    the run finishes with no provenance at all, which is the case this file
+    exists to prevent.
+    """
+    from pyhealth.utils import write_run_config
+
+    config = {
+        "output_dir": Path("/scratch/run"),
+        "device": torch.device("cpu"),
+        "window_hours": 24,
+        "model": nn.Linear(2, 2),
+    }
+    record = json.loads(Path(write_run_config(str(tmp_path), config)).read_text())
+
+    assert set(record["config"]) == set(config)
+    assert record["config"]["window_hours"] == 24
+
+
+def test_the_temporary_file_does_not_remain_after_a_write(tmp_path):
+    """A partially written run_config.json is worse than none, so the write is
+    atomic. No temporary file may survive it.
+    """
+    from pyhealth.utils import write_run_config
+
+    write_run_config(str(tmp_path), {"seed": 1})
+
+    leftovers = [p.name for p in tmp_path.iterdir() if ".tmp." in p.name]
+    assert leftovers == [], f"temporary files remained: {leftovers}"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Discriminative learning rate for the text pathway
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+TEXT_FIELDS = {"notes"}
+
+
+def test_a_trainable_encoder_keeps_the_projection_at_the_base_rate():
+    """``encoder_lr`` exists to give a PRETRAINED encoder a gentler rate. A
+    projection with random values must not receive that rate.
+    """
+    from pyhealth.trainer import is_text_pathway_param
+
+    assert is_text_pathway_param(
+        "embedding_model.encoders.notes.layer.0.weight", TEXT_FIELDS
+    )
+    assert not is_text_pathway_param(
+        "embedding_model.projections.notes.weight", TEXT_FIELDS
+    )
+
+
+def test_a_frozen_encoder_puts_the_projection_in_the_group():
+    """With the encoder frozen, every ``encoders.*`` parameter has
+    ``requires_grad=False``, so the group is empty and ``encoder_lr`` controls
+    nothing. The projection is then the only trainable text parameter, so it IS
+    the text pathway.
+    """
+    from pyhealth.trainer import is_text_pathway_param
+
+    assert is_text_pathway_param(
+        "embedding_model.projections.notes.weight", TEXT_FIELDS, frozen_text_fields={"notes"}
+    )
+
+
+def test_a_non_text_parameter_never_joins_the_group():
+    from pyhealth.trainer import is_text_pathway_param
+
+    for name in (
+        "_unified_backbone.layers.0.weight",
+        "fc.weight",
+        "embedding_model.encoders.labs.weight",
+        "embedding_model.projections.labs.weight",
+    ):
+        assert not is_text_pathway_param(name, TEXT_FIELDS, frozen_text_fields={"notes"})
+
+
+def test_a_field_name_that_is_a_prefix_of_another_does_not_match():
+    """``notes`` must not capture ``notes_extra``."""
+    from pyhealth.trainer import is_text_pathway_param
+
+    assert not is_text_pathway_param(
+        "embedding_model.encoders.notes_extra.weight", TEXT_FIELDS
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Within-epoch loss trajectory
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def test_the_epoch_record_holds_the_within_epoch_trajectory():
+    """An epoch mean cannot show the difference between a run that starts badly
+    and a run that becomes worse inside the epoch.
+    """
+    import inspect
+
+    from pyhealth.trainer import Trainer
+
+    source = inspect.getsource(Trainer.train)
+    for field in (
+        "train_loss_first_step",
+        "train_loss_first100",
+        "train_loss_last100",
+    ):
+        assert f'"{field}"' in source, f"{field} is not recorded"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Chest X-ray as a third modality
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+CXR_ARMS = [
+    ("cxr_only", {}, {"cxr"}),
+    ("cxr_labs", {"include_labs": True}, {"cxr", "labs", "labs_mask"}),
+    (
+        "cxr_notes_labs",
+        {"include_labs": True, "include_notes": True},
+        {"cxr", "labs", "labs_mask", "admission_note_times"},
+    ),
+]
+
+
+@pytest.mark.parametrize("name,kwargs,expected", CXR_ARMS)
+def test_each_cxr_arm_declares_the_fields_it_uses(name, kwargs, expected):
+    from pyhealth.tasks import CXRMultimodalMIMIC4
+
+    task = CXRMultimodalMIMIC4(window_hours=24, **kwargs)
+
+    assert set(task.input_schema) == expected
+
+
+def test_the_lab_mask_matches_the_other_tasks():
+    """``StageNetTensorProcessor`` takes no arguments. A processor option here
+    that the other tasks do not use would raise at dataset build time.
+    """
+    from pyhealth.processors import StageNetTensorProcessor
+    from pyhealth.tasks import CXRMultimodalMIMIC4
+
+    task = CXRMultimodalMIMIC4(window_hours=24, include_labs=True)
+    name, options = task.input_schema["labs_mask"]
+
+    assert name == "stagenet_tensor"
+    assert options == {}
+    StageNetTensorProcessor(**options)  # must not raise
+
+
+def test_an_image_uses_the_same_time_convention_as_a_laboratory_value():
+    """``StudyDate`` and ``StudyTime`` give hours from admission, which is what
+    the unified embedding reads for every modality.
+    """
+    from pyhealth.tasks import CXRMultimodalMIMIC4
+
+    task = CXRMultimodalMIMIC4(window_hours=24)
+    assert task.input_schema["cxr"][0] == "time_image"

From 50a8a873d39409c37a8814e42940d698da4a7d80 Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Wed, 12 Aug 2026 13:56:19 -0400
Subject: [PATCH 02/12] Include the task in the run directory name

The run directory was named from the model and the seed only. A paired
comparison holds both fixed and varies the task, so --task labs_only and
--task notes_labs at seed 42 resolved to one directory: transformer_seed42.
The second run overwrote the first run's metrics_history.json, run_config.json
and predictions CSV.

The loss is silent. The surviving directory looks like a complete run, and the
provenance this PR adds would describe only the arm that finished last.

Found while reviewing the same defect in the upstream consolidation branch.
---
 .../unified_embedding_e2e_mimic4.py           |  8 ++-
 tests/test_run_directory_naming.py            | 72 +++++++++++++++++++
 2 files changed, 78 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_run_directory_naming.py

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 186b2998a..d3ed489e8 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -655,8 +655,12 @@ def run(args: argparse.Namespace) -> Path:
             RuntimeWarning, stacklevel=2,
         )
 
-    # Experiment name encodes model + seed for easy log separation
-    exp_name = f"{args.model}_seed{args.seed}"
+    # The task MUST be in the name. Without it, two arms of the same comparison
+    # at the same seed, for example --task labs_only and --task notes_labs,
+    # resolve to one directory and the second run overwrites the first run's
+    # metrics_history.json, run_config.json and predictions CSV. The loss is
+    # silent: the surviving directory looks like a complete run.
+    exp_name = f"{args.task}_{args.model}_seed{args.seed}"
     output_dir = Path(args.output_dir)
 
     trainer = Trainer(
diff --git a/tests/test_run_directory_naming.py b/tests/test_run_directory_naming.py
new file mode 100644
index 000000000..8d8475ab6
--- /dev/null
+++ b/tests/test_run_directory_naming.py
@@ -0,0 +1,72 @@
+"""Two arms of one comparison must not write to the same directory.
+
+The run directory was named from the model and the seed only. A paired
+comparison holds both of those fixed and varies the task, so the two arms
+resolved to one path and the second run overwrote the first run's
+``metrics_history.json``, ``run_config.json`` and predictions CSV. The loss is
+silent: the surviving directory looks like a complete run.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+RUNNER = (
+    Path(__file__).resolve().parents[1]
+    / "examples"
+    / "mortality_prediction"
+    / "unified_embedding_e2e_mimic4.py"
+)
+
+
+def _exp_name_template() -> str:
+    """The literal assignment, read from source.
+
+    The name is computed deep inside ``main()`` after a dataset build, so the
+    assignment itself is what this test pins.
+    """
+    for line in RUNNER.read_text().splitlines():
+        stripped = line.strip()
+        if stripped.startswith("exp_name = "):
+            return stripped
+    raise AssertionError("exp_name is no longer assigned in the runner")
+
+
+def _render(task: str, model: str, seed: int) -> str:
+    template = _exp_name_template().split("=", 1)[1].strip()
+    args = SimpleNamespace(task=task, model=model, seed=seed)
+    return eval(template, {"args": args})  # noqa: S307 - a literal f-string
+
+
+def test_the_run_directory_name_includes_the_task():
+    assert "args.task" in _exp_name_template(), (
+        "two arms of the same comparison would share one output directory"
+    )
+
+
+@pytest.mark.parametrize(
+    "left,right",
+    [
+        (("labs_only", "transformer", 42), ("notes_labs", "transformer", 42)),
+        (("labs_only", "rnn", 1), ("cxr_notes_labs", "rnn", 1)),
+    ],
+)
+def test_two_arms_at_the_same_seed_get_different_directories(left, right):
+    assert _render(*left) != _render(*right)
+
+
+def test_the_seed_still_separates_repeats_of_one_arm():
+    assert _render("notes_labs", "transformer", 42) != _render(
+        "notes_labs", "transformer", 43
+    )
+
+
+def test_the_model_still_separates_backbones():
+    assert _render("notes_labs", "transformer", 42) != _render(
+        "notes_labs", "ehrmamba", 42
+    )

From 4215fdcf2048de3f145076a5bcdc857424b347fb Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Wed, 12 Aug 2026 14:52:30 -0400
Subject: [PATCH 03/12] Match the task-class API in the runner, and refuse to
 drop a flag silently

The runner called NotesLabsMIMIC4 with include_labs, note_extraction,
note_source, discharge_note_policy and text_normalize. The task class on this
branch accepts none of them, so --task notes_labs, the primary arm of the
comparison in this PR, died at construction:

  TypeError: NotesLabsMIMIC4.__init__() got an unexpected keyword argument
  'include_labs'

The runner now passes only the parameters the task class declares. A flag that
the class cannot honour stops the run instead of being dropped, because a
silently ignored --discharge-note-policy would record one protocol in
run_config.json and execute another, which is the failure mode this PR exists
to remove.

Note collection on this branch therefore uses the section extraction that main
already has, through _collect_notes with DISCHARGE_CLINICAL_HEADERS.

Caught by a smoke run on real MIMIC-IV rather than by the unit tests, which do
not build a task from parsed arguments.
---
 .../unified_embedding_e2e_mimic4.py           | 42 ++++++++++++++++---
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index d3ed489e8..90faff630 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -50,6 +50,7 @@
 from __future__ import annotations
 
 import argparse
+import inspect
 import csv
 import warnings
 from pathlib import Path
@@ -149,7 +150,11 @@ def _build_task(args: argparse.Namespace):
     if args.task == "clinical_notes_icd_labs":
         return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours)
     if args.task in ("notes_labs", "notes_only"):
-        task_kwargs = dict(
+        # Pass only what this checkout's task class accepts. A flag the class
+        # cannot honour must stop the run, not be dropped: a silently ignored
+        # --discharge-note-policy would record one protocol in run_config.json
+        # and execute another.
+        wanted = dict(
             window_hours=args.observation_window_hours,
             include_icd=args.icd_codes,
             include_vitals=args.include_vitals,
@@ -158,12 +163,37 @@ def _build_task(args: argparse.Namespace):
             note_source=getattr(args, "note_source", "discharge"),
             discharge_note_policy=getattr(
                 args, "discharge_note_policy", "extraction"),
+            text_normalize=getattr(args, "text_normalize", "none"),
         )
-        # Only pass text_normalize when actually requested, so this script still
-        # runs against a checkout whose task class predates the parameter.
-        _tn = getattr(args, "text_normalize", "none")
-        if _tn and _tn != "none":
-            task_kwargs["text_normalize"] = _tn
+        accepted = set(
+            inspect.signature(NotesLabsMIMIC4.__init__).parameters
+        )
+        defaults = {
+            "include_icd": False,
+            "include_vitals": False,
+            "include_labs": True,
+            "note_extraction": "regex",
+            "note_source": "discharge",
+            "discharge_note_policy": "extraction",
+            "text_normalize": "none",
+        }
+        unsupported = [
+            name
+            for name, value in wanted.items()
+            if name not in accepted and value != defaults.get(name)
+        ]
+        if unsupported:
+            raise SystemExit(
+                f"NotesLabsMIMIC4 in this checkout does not accept "
+                f"{', '.join(sorted(unsupported))}. Either drop the flag or use "
+                f"a checkout whose task class supports it."
+            )
+        task_kwargs = {k: v for k, v in wanted.items() if k in accepted}
+        if args.task == "notes_only" and "include_labs" not in accepted:
+            raise SystemExit(
+                "--task notes_only needs a NotesLabsMIMIC4 that accepts "
+                "include_labs; this checkout always emits labs."
+            )
         task = NotesLabsMIMIC4(**task_kwargs)
         if args.tokenizer_model:
             schema_key = "admission_note_times"

From 9caa81ca7647147be3ffe17a00766564a6a228af Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Wed, 12 Aug 2026 15:18:15 -0400
Subject: [PATCH 04/12] Call the branch API in the runner: freeze_text_encoder,
 and AMP on train()

Three more calls were written against a checkout this branch does not have.

UnifiedMultimodalEmbeddingModel takes freeze_text_encoder, a boolean, not
text_finetune_mode, so every run died at model construction:

  TypeError: UnifiedMultimodalEmbeddingModel.__init__() got an unexpected
  keyword argument 'text_finetune_mode'

use_amp and amp_dtype are parameters of Trainer.train, not of Trainer, because
mixed precision is a property of the training loop and not of the object.

get_dataloader gains num_workers in the performance PR, which is not in this
branch, so the audit batch uses the default loader.

An AST check over the runner now reports no keyword argument that the imported
pyhealth signatures reject.
---
 .../mortality_prediction/unified_embedding_e2e_mimic4.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 90faff630..8cd40b16b 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -288,7 +288,7 @@ def _build_model(
         processors=sample_dataset.input_processors,
         embedding_dim=args.embedding_dim,
         field_embeddings=field_embeddings,
-        text_finetune_mode=finetune_mode,
+        freeze_text_encoder=(finetune_mode == "frozen"),
         normalize_content=not getattr(args, "no_normalize_content", False),
         numeric_standardizers=numeric_standardizers,
         cache_frozen_text=not getattr(args, "no_text_cache", False),
@@ -609,7 +609,6 @@ def run(args: argparse.Namespace) -> Path:
         # misleading difference caused only by independent shuffling.
         audit_batch = next(iter(get_dataloader(
             train_ds, batch_size=args.batch_size, shuffle=False,
-            num_workers=args.loader_num_workers,
         )))
         with torch.no_grad():
             model(**audit_batch)
@@ -700,8 +699,6 @@ def run(args: argparse.Namespace) -> Path:
         enable_logging=True,
         output_path=str(output_dir),
         exp_name=exp_name,
-        use_amp=args.use_amp,
-        amp_dtype=args.amp_dtype,
     )
 
     # BottleneckTransformer is more fragile on full MIMIC-IV with no warmup.
@@ -774,6 +771,10 @@ def run(args: argparse.Namespace) -> Path:
             load_best_model_at_last=True,
             patience=args.patience,
             encoder_lr=args.encoder_lr,
+            # Mixed precision is a property of the training loop, not of the
+            # Trainer object.
+            use_amp=args.use_amp,
+            amp_dtype=args.amp_dtype,
         )
 
     y_true, y_prob, _, patient_ids = trainer.inference(

From 66236233955c4b1da7725421e76f35392ac1056c Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Wed, 12 Aug 2026 15:28:29 -0400
Subject: [PATCH 05/12] Filter DataLoader options to the installed
 get_dataloader signature

The worker options were expanded into get_dataloader with **loader_kwargs,
which an AST check over keyword arguments cannot see, so the previous audit
reported the runner clean while every run still died:

  TypeError: get_dataloader() got an unexpected keyword argument 'num_workers'

Those options arrive with the performance PR, which is not in this branch. The
runner now passes only what the installed signature declares, and stops if the
caller asked for an option it cannot honour, so a requested option is never
silently ignored.
---
 .../unified_embedding_e2e_mimic4.py           | 26 ++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 8cd40b16b..d23ea3e7d 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -626,7 +626,10 @@ def run(args: argparse.Namespace) -> Path:
     print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.")
     model._pos_weight = torch.tensor([pw_value], dtype=torch.float32)
 
-    loader_kwargs = {
+    # DataLoader worker options arrive with the performance PR. Pass only what
+    # this checkout's get_dataloader accepts, and stop if the caller asked for
+    # one it cannot honour, so a requested option is never silently ignored.
+    _wanted_loader = {
         "num_workers": args.loader_num_workers,
         "pin_memory": args.pin_memory,
         "persistent_workers": (
@@ -636,6 +639,27 @@ def run(args: argparse.Namespace) -> Path:
             args.prefetch_factor if args.loader_num_workers > 0 else None
         ),
     }
+    _loader_accepts = set(inspect.signature(get_dataloader).parameters)
+    _loader_defaults = {
+        "num_workers": 0,
+        "pin_memory": False,
+        "persistent_workers": False,
+        "prefetch_factor": None,
+    }
+    _unsupported_loader = [
+        name
+        for name, value in _wanted_loader.items()
+        if name not in _loader_accepts and value != _loader_defaults[name]
+    ]
+    if _unsupported_loader:
+        raise SystemExit(
+            f"get_dataloader in this checkout does not accept "
+            f"{', '.join(sorted(_unsupported_loader))}. Drop the flag, or use a "
+            f"checkout that includes the DataLoader worker options."
+        )
+    loader_kwargs = {
+        k: v for k, v in _wanted_loader.items() if k in _loader_accepts
+    }
     train_loader = get_dataloader(
         train_ds,
         batch_size=args.batch_size,

From 83126a38eecc2a47833c1e87b185930c53dcf752 Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Wed, 12 Aug 2026 15:38:43 -0400
Subject: [PATCH 06/12] Request only the metrics this checkout implements

binary_metrics_fn has no f1_opt, so validation aborted at the end of epoch 1:

  ValueError: Unknown metric for binary classification: f1_opt

Model selection uses pr_auc, a rank metric that needs no threshold, so dropping
the threshold-optimised F1 does not change which checkpoint is chosen.
---
 .../mortality_prediction/unified_embedding_e2e_mimic4.py    | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index d23ea3e7d..020bc68e5 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -718,7 +718,11 @@ def run(args: argparse.Namespace) -> Path:
 
     trainer = Trainer(
         model=model,
-        metrics=["pr_auc", "roc_auc", "f1", "f1_opt", "accuracy"],
+        # f1_opt is a threshold-optimised F1 that this checkout's
+        # binary_metrics_fn does not implement, and requesting it aborts
+        # validation at the end of epoch 1. Model selection uses pr_auc, which
+        # is a rank metric and needs no threshold.
+        metrics=["pr_auc", "roc_auc", "f1", "accuracy"],
         device=args.device,
         enable_logging=True,
         output_path=str(output_dir),

From ea22ff9d137a0dce1e0afbf4a4309cf725b5828c Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Thu, 13 Aug 2026 10:53:28 -0400
Subject: [PATCH 07/12] Expose the CXR layout variant in the runner

MIMIC4Dataset takes cxr_variant, and the runner never passed it, so every CXR
run used the default layout. The resized set that this project uses is the
sunlab layout, and the default config expects a column it does not have, so all
three CXR arms failed at dataset build:

  KeyError: 'studytime_normalized'

The default config reads mimic-cxr-2.0.0-metadata-pyhealth.csv and needs
studytime_normalized. The sunlab variant reads the resized set, normalises
StudyTime itself, and derives image paths from dicom_id.

The help text names the exact failure so the wrong choice is diagnosable from
the flag rather than from a pandas KeyError inside dask.
---
 .../unified_embedding_e2e_mimic4.py                | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 020bc68e5..0637e64be 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -136,6 +136,7 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
         note_tables=note_tables,
         cxr_root=args.cxr_root if cxr_tables else None,
         cxr_tables=cxr_tables,
+        cxr_variant=args.cxr_variant,
         cache_dir=args.cache_dir,
         dev=args.dev if args.dev else False,
         num_workers=args.num_workers,
@@ -821,6 +822,19 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--ehr-root", type=str, required=True)
     parser.add_argument("--note-root", type=str, default=None)
     parser.add_argument("--cxr-root", type=str, default=None)
+    parser.add_argument(
+        "--cxr-variant",
+        choices=["default", "sunlab"],
+        default="default",
+        help=(
+            "Layout of the CXR root. 'default' reads "
+            "mimic-cxr-2.0.0-metadata-pyhealth.csv and needs a "
+            "studytime_normalized column. 'sunlab' reads the resized set, "
+            "normalises StudyTime itself and derives image paths from "
+            "dicom_id. Choosing the wrong one fails at dataset build with "
+            "KeyError: 'studytime_normalized'."
+        ),
+    )
     parser.add_argument("--cache-dir", type=str, default=None)
     parser.add_argument("--output-dir", type=str, default="./output/unified_e2e")
     parser.add_argument(

From e8853e6306fba5e804885a6edaebdb94d1cd909b Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Thu, 13 Aug 2026 10:56:07 -0400
Subject: [PATCH 08/12] Resolve the flattened CXR image directory instead of
 hardcoding one name

The sunlab CXR variant required a directory literally named "images". The
resized set this project uses holds the same flattened {dicom_id}.jpg files
under "resized_images", so all three CXR arms failed on a complete and correct
dataset of 377,110 images:

  FileNotFoundError: Sunlab images directory not found: .../images

Both names are now accepted, the derived image_path follows whichever was
found, and the error lists what was looked for.
---
 pyhealth/datasets/mimic4.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py
index 9be17e1b8..a4e1a72fa 100644
--- a/pyhealth/datasets/mimic4.py
+++ b/pyhealth/datasets/mimic4.py
@@ -282,12 +282,25 @@ def prepare_metadata(self, root: str) -> None:
                 "Expected existing metadata linked by dicom_id/subject_id/study_id."
             )
 
-        images_dir = os.path.join(root, "images")
-        if not os.path.isdir(images_dir):
+        # The flattened layout appears under more than one directory name
+        # depending on how the set was produced, so accept either rather than
+        # hardcoding one and failing on a complete, correct dataset.
+        candidates = ("images", "resized_images")
+        images_dir = next(
+            (
+                os.path.join(root, name)
+                for name in candidates
+                if os.path.isdir(os.path.join(root, name))
+            ),
+            None,
+        )
+        if images_dir is None:
             raise FileNotFoundError(
-                f"Sunlab images directory not found: {images_dir}. "
-                "Expected flattened image files at images/{dicom_id}.jpg."
+                f"No flattened image directory under {root}. Looked for "
+                f"{', '.join(candidates)}, each expected to hold "
+                "{dicom_id}.jpg."
             )
+        images_subdir = os.path.basename(images_dir)
 
         metadata = pd.read_csv(metadata_path, dtype=str)
 
@@ -314,7 +327,7 @@ def normalize_studytime(value: Optional[str]) -> str:
         metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime)
 
         metadata["image_path"] = metadata[dicom_col].apply(
-            lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg")
+            lambda dicom_id: os.path.join(root, images_subdir, f"{dicom_id}.jpg")
         )
 
         # Align with existing config conventions by using lowercase headers.

From b03a1723a3125d3cccc741e72a403b44a7b3ee2a Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Thu, 13 Aug 2026 12:12:10 -0400
Subject: [PATCH 09/12] Test that the sunlab CXR layout accepts resized_images

The last commit accepted both directory names after a complete cohort failed
on a hardcoded images path. The class docstring still named only images, and
no test covered the lookup.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 pyhealth/datasets/mimic4.py               |  2 +-
 tests/test_run_provenance_and_pathways.py | 43 +++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py
index a4e1a72fa..e569a34b1 100644
--- a/pyhealth/datasets/mimic4.py
+++ b/pyhealth/datasets/mimic4.py
@@ -235,7 +235,7 @@ class MIMIC4CXRSunlabDataset(BaseDataset):
     Sunlab variant of the MIMIC-CXR Chest X-ray dataset.
 
     This variant uses the existing metadata CSV and derives flattened image
-    paths at ``images/{dicom_id}.jpg``.
+    paths at ``{images|resized_images}/{dicom_id}.jpg``.
     """
 
     def __init__(
diff --git a/tests/test_run_provenance_and_pathways.py b/tests/test_run_provenance_and_pathways.py
index 6d686a818..f245bb269 100644
--- a/tests/test_run_provenance_and_pathways.py
+++ b/tests/test_run_provenance_and_pathways.py
@@ -213,3 +213,46 @@ def test_an_image_uses_the_same_time_convention_as_a_laboratory_value():
 
     task = CXRMultimodalMIMIC4(window_hours=24)
     assert task.input_schema["cxr"][0] == "time_image"
+
+
+def test_sunlab_accepts_resized_images_as_well_as_images(tmp_path):
+    """The resized set holds flattened ``{dicom_id}.jpg`` files under
+    ``resized_images``. Requiring a directory literally named ``images`` failed
+    on a complete 377,110-image cohort.
+    """
+    import pandas as pd
+    from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset
+
+    root = tmp_path / "cxr"
+    (root / "resized_images").mkdir(parents=True)
+    pd.DataFrame(
+        {"dicom_id": ["abc123"], "StudyTime": ["123045.0"]}
+    ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False)
+
+    MIMIC4CXRSunlabDataset.prepare_metadata(
+        object.__new__(MIMIC4CXRSunlabDataset), str(root)
+    )
+
+    out = pd.read_csv(root / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv")
+    path = str(out["image_path"].iloc[0])
+    assert path.endswith(f"resized_images{os.sep}abc123.jpg") or path.endswith(
+        "resized_images/abc123.jpg"
+    )
+    assert str(out["studytime"].iloc[0]).zfill(6) == "123045"
+
+
+def test_sunlab_reports_both_directory_names_when_neither_exists(tmp_path):
+    import pandas as pd
+    from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset
+
+    root = tmp_path / "cxr"
+    root.mkdir()
+    pd.DataFrame(
+        {"dicom_id": ["abc123"], "StudyTime": ["1"]}
+    ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False)
+
+    with pytest.raises(FileNotFoundError, match="resized_images"):
+        MIMIC4CXRSunlabDataset.prepare_metadata(
+            object.__new__(MIMIC4CXRSunlabDataset), str(root)
+        )
+

From 330623409ee997eed382f00e20d7631a43f1784d Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Thu, 13 Aug 2026 14:10:48 -0400
Subject: [PATCH 10/12] Derive the image channel count from the processor

The unified embedding sizes its patch embedding from processor.in_channels and
falls back to 3. TimeImageProcessor never exposed that attribute, so a greyscale
CXR task built a 3-channel patch embedding and fed it 1-channel images. The
mismatch appeared only at the first forward pass, after the full image cache had
been built:

  RuntimeError: Given groups=1, weight of size [128, 3, 16, 16],
  expected input[16, 1, 224, 224] to have 3 channels, but got 1

in_channels now follows n_channels when set and otherwise the PIL mode, matching
_zero_image_tensor exactly so a placeholder cannot differ from a real image.
---
 pyhealth/processors/time_image_processor.py | 20 ++++++++++++++++++
 tests/test_run_provenance_and_pathways.py   | 23 +++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py
index 421998d07..a5cc4fafc 100644
--- a/pyhealth/processors/time_image_processor.py
+++ b/pyhealth/processors/time_image_processor.py
@@ -178,6 +178,26 @@ def _zero_image_tensor(self) -> torch.Tensor:
             c = 3
         return torch.zeros(c, self.image_size, self.image_size)
 
+    @property
+    def in_channels(self) -> int:
+        """Channel count implied by ``mode``.
+
+        The unified embedding sizes its patch embedding from this. Without it
+        the model defaulted to 3 while a greyscale task produced 1, and the
+        mismatch only appeared at the first forward pass:
+
+            RuntimeError: Given groups=1, weight of size [128, 3, 16, 16],
+            expected input[16, 1, 224, 224] to have 3 channels
+
+        Deriving it here means the two cannot disagree.
+        """
+        # Must match _zero_image_tensor exactly, or a placeholder image would
+        # carry a different channel count from a real one.
+        if self.n_channels is not None:
+            return int(self.n_channels)
+        return {"1": 1, "L": 1, "LA": 2, "RGB": 3, "RGBA": 4}.get(self.mode or "RGB", 3)
+
+
     def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor:
         """Load and transform a single image from disk.
 
diff --git a/tests/test_run_provenance_and_pathways.py b/tests/test_run_provenance_and_pathways.py
index f245bb269..d0368b226 100644
--- a/tests/test_run_provenance_and_pathways.py
+++ b/tests/test_run_provenance_and_pathways.py
@@ -256,3 +256,26 @@ def test_sunlab_reports_both_directory_names_when_neither_exists(tmp_path):
             object.__new__(MIMIC4CXRSunlabDataset), str(root)
         )
 
+
+
+def test_the_image_channel_count_comes_from_the_processor():
+    """The unified embedding sizes its patch embedding from
+    ``processor.in_channels``. Without that attribute it fell back to 3 while a
+    greyscale CXR task produced 1, and the mismatch surfaced only at the first
+    forward pass:
+
+        RuntimeError: Given groups=1, weight of size [128, 3, 16, 16],
+        expected input[16, 1, 224, 224] to have 3 channels
+    """
+    from pyhealth.processors import TimeImageProcessor
+
+    assert TimeImageProcessor(image_size=224, mode="L").in_channels == 1
+    assert TimeImageProcessor(image_size=224, mode="RGB").in_channels == 3
+
+
+def test_a_placeholder_image_has_the_same_channels_as_a_real_one():
+    from pyhealth.processors import TimeImageProcessor
+
+    for mode in ("L", "RGB"):
+        processor = TimeImageProcessor(image_size=64, mode=mode)
+        assert processor._zero_image_tensor().shape[0] == processor.in_channels

From a320943dddc68c6ecd3880c8873ee4ac560587ac Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Mon, 17 Aug 2026 21:21:03 -0400
Subject: [PATCH 11/12] Write sunlab CXR metadata to cache when the data root
 is read-only.

prepare_metadata wrote mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv into the PhysioNet root. That path is not writable on the cluster, so CXR setup failed after a complete image directory had already been found. Cache is tried first, and the generated YAML is rewritten to the absolute CSV path. The test chmods the root to 555 and checks the CSV lands in cache.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 pyhealth/datasets/mimic4.py               | 53 ++++++++++++++++++++---
 tests/test_run_provenance_and_pathways.py | 34 +++++++++++++++
 2 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py
index e569a34b1..db47697b4 100644
--- a/pyhealth/datasets/mimic4.py
+++ b/pyhealth/datasets/mimic4.py
@@ -252,7 +252,11 @@ def __init__(
                 os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml"
             )
             logger.info(f"Using default Sunlab CXR config: {config_path}")
-        self.prepare_metadata(root)
+        metadata_csv = self.prepare_metadata(root, cache_dir=cache_dir)
+        if os.path.dirname(os.path.abspath(metadata_csv)) != os.path.abspath(root):
+            config_path = self._rewrite_sunlab_config(
+                config_path, metadata_csv, cache_dir or os.path.dirname(metadata_csv)
+            )
         log_memory_usage(f"Before initializing {dataset_name}")
         super().__init__(
             root=root,
@@ -274,7 +278,26 @@ def _resolve_column_name(columns: List[str], target: str) -> str:
             )
         return resolved
 
-    def prepare_metadata(self, root: str) -> None:
+    @staticmethod
+    def _rewrite_sunlab_config(
+        config_path: str, metadata_csv: str, dest_dir: str
+    ) -> str:
+        """Point the sunlab YAML at a metadata CSV that is not under root."""
+        with open(config_path, encoding="utf-8") as f:
+            text = f.read()
+        rewritten = text.replace(
+            "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv",
+            metadata_csv,
+        )
+        os.makedirs(dest_dir, exist_ok=True)
+        out = os.path.join(dest_dir, "mimic4_cxr_sunlab.generated.yaml")
+        with open(out, "w", encoding="utf-8") as f:
+            f.write(rewritten)
+        return out
+
+    def prepare_metadata(
+        self, root: str, cache_dir: Optional[str] = None
+    ) -> str:
         metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv")
         if not os.path.exists(metadata_path):
             raise FileNotFoundError(
@@ -333,9 +356,29 @@ def normalize_studytime(value: Optional[str]) -> str:
         # Align with existing config conventions by using lowercase headers.
         metadata.columns = [col.lower() for col in metadata.columns]
 
-        metadata.to_csv(
-            os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"),
-            index=False,
+        filename = "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+        dest_dirs = []
+        if cache_dir:
+            dest_dirs.append(str(cache_dir))
+        dest_dirs.append(root)
+
+        for d in dest_dirs:
+            existing = os.path.join(d, filename)
+            if os.path.isfile(existing):
+                return existing
+
+        last_err: Optional[OSError] = None
+        for d in dest_dirs:
+            os.makedirs(d, exist_ok=True)
+            dest = os.path.join(d, filename)
+            try:
+                metadata.to_csv(dest, index=False)
+                return dest
+            except OSError as exc:
+                last_err = exc
+                continue
+        raise PermissionError(
+            f"Could not write {filename} under {dest_dirs}: {last_err}"
         )
 
 
diff --git a/tests/test_run_provenance_and_pathways.py b/tests/test_run_provenance_and_pathways.py
index d0368b226..5302e805c 100644
--- a/tests/test_run_provenance_and_pathways.py
+++ b/tests/test_run_provenance_and_pathways.py
@@ -241,6 +241,40 @@ def test_sunlab_accepts_resized_images_as_well_as_images(tmp_path):
     assert str(out["studytime"].iloc[0]).zfill(6) == "123045"
 
 
+def test_sunlab_writes_metadata_to_cache_when_root_is_unwritable(tmp_path):
+    """PhysioNet roots are typically read-only. Writing the derived CSV there
+    raised PermissionError after a complete cohort had already been found.
+    """
+    import pandas as pd
+    from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset
+
+    root = tmp_path / "cxr"
+    (root / "resized_images").mkdir(parents=True)
+    pd.DataFrame(
+        {"dicom_id": ["abc"], "StudyTime": ["93000"], "subject_id": ["1"]}
+    ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False)
+    cache = tmp_path / "cache"
+    cache.mkdir()
+    os.chmod(root, 0o555)
+    try:
+        dest = MIMIC4CXRSunlabDataset.prepare_metadata(
+            object.__new__(MIMIC4CXRSunlabDataset),
+            str(root),
+            cache_dir=str(cache),
+        )
+    finally:
+        os.chmod(root, 0o755)
+
+    assert dest.startswith(str(cache))
+    assert Path(dest).is_file()
+    assert not (root / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv").exists()
+    written = pd.read_csv(dest)
+    path = str(written.loc[0, "image_path"])
+    assert path.endswith(f"resized_images{os.sep}abc.jpg") or path.endswith(
+        "resized_images/abc.jpg"
+    )
+
+
 def test_sunlab_reports_both_directory_names_when_neither_exists(tmp_path):
     import pandas as pd
     from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset

From 7776fbefaa5e9d2323c4ff9d262b5be83fa979a2 Mon Sep 17 00:00:00 2001
From: Rian354 <rianrian290@icloud.com>
Date: Mon, 17 Aug 2026 21:21:12 -0400
Subject: [PATCH 12/12] Default JambaEHR to 2 transformer + 2 mamba layers.

The runner compared Jamba against Transformer and RNN while the library defaulted to 2+6, so the extra Mamba stack was an uncontrolled capacity difference. Class and CLI defaults are now 2+2; the test checks both the constructor defaults and the layer schedule.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../unified_embedding_e2e_mimic4.py           |  4 +-
 pyhealth/models/jamba_ehr.py                  |  8 ++--
 tests/test_jamba_default_depth.py             | 43 +++++++++++++++++++
 3 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_jamba_default_depth.py

diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
index 0637e64be..e24a9d036 100644
--- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -1184,9 +1184,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--jamba-transformer-layers", type=int, default=2,
                         help="Number of Transformer (attention) layers in JambaEHR. "
                              "Standard: 1 (a single Jamba block = 1 attn + 1 mamba).")
-    parser.add_argument("--jamba-mamba-layers", type=int, default=6,
+    parser.add_argument("--jamba-mamba-layers", type=int, default=2,
                         help="Number of Mamba (SSM) layers in JambaEHR. "
-                             "Standard: 1 (a single Jamba block = 1 attn + 1 mamba).")
+                             "Default 2, matching the library JambaEHR default.")
 
     return parser.parse_args()
 
diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py
index 37d738f3f..7b8062eae 100644
--- a/pyhealth/models/jamba_ehr.py
+++ b/pyhealth/models/jamba_ehr.py
@@ -73,7 +73,7 @@ class JambaLayer(nn.Module):
     Args:
         feature_size (int): Hidden dimension shared by all layers.
         num_transformer_layers (int): Number of attention layers. Default 2.
-        num_mamba_layers (int): Number of SSM layers. Default 6.
+        num_mamba_layers (int): Number of SSM layers. Default 2.
         heads (int): Attention heads for Transformer layers. Default 4.
         dropout (float): Dropout rate for Transformer layers. Default 0.3.
         state_size (int): SSM state size for Mamba layers. Default 16.
@@ -94,7 +94,7 @@ def __init__(
         self,
         feature_size: int,
         num_transformer_layers: int = 2,
-        num_mamba_layers: int = 6,
+        num_mamba_layers: int = 2,
         heads: int = 4,
         dropout: float = 0.3,
         state_size: int = 16,
@@ -187,7 +187,7 @@ class JambaEHR(BaseModel):
         dataset (SampleDataset): Dataset providing processed inputs.
         embedding_dim (int): Embedding and hidden dimension. Default 128.
         num_transformer_layers (int): Transformer layers per stream. Default 2.
-        num_mamba_layers (int): Mamba layers per stream. Default 6.
+        num_mamba_layers (int): Mamba layers per stream. Default 2.
         heads (int): Attention heads per Transformer block. Default 4.
         dropout (float): Dropout rate. Default 0.3.
         state_size (int): SSM state size in Mamba blocks. Default 16.
@@ -237,7 +237,7 @@ def __init__(
         dataset: SampleDataset,
         embedding_dim: int = 128,
         num_transformer_layers: int = 2,
-        num_mamba_layers: int = 6,
+        num_mamba_layers: int = 2,
         heads: int = 4,
         dropout: float = 0.3,
         state_size: int = 16,
diff --git a/tests/test_jamba_default_depth.py b/tests/test_jamba_default_depth.py
new file mode 100644
index 000000000..5951d60fe
--- /dev/null
+++ b/tests/test_jamba_default_depth.py
@@ -0,0 +1,43 @@
+"""Proof that JambaEHR defaults to 2 transformer + 2 mamba layers.
+
+The published runner used library defaults of 2+6, which over-parameterised
+the e2e comparison against Transformer and RNN. Both the class and the CLI
+now default to 2+2.
+"""
+
+from __future__ import annotations
+
+import re
+import unittest
+from pathlib import Path
+
+RUNNER = (
+    Path(__file__).resolve().parents[1]
+    / "examples"
+    / "mortality_prediction"
+    / "unified_embedding_e2e_mimic4.py"
+)
+
+
+class TestJambaDefaultDepth(unittest.TestCase):
+    def test_library_defaults_to_two_plus_two(self):
+        from pyhealth.models.jamba_ehr import JambaEHR, JambaLayer
+
+        self.assertEqual(JambaLayer.__init__.__defaults__[0], 2)
+        self.assertEqual(JambaLayer.__init__.__defaults__[1], 2)
+        self.assertEqual(JambaEHR.__init__.__defaults__[1], 2)
+        self.assertEqual(JambaEHR.__init__.__defaults__[2], 2)
+        layer = JambaLayer(16)
+        self.assertEqual(layer.schedule.count("transformer"), 2)
+        self.assertEqual(layer.schedule.count("mamba"), 2)
+
+    def test_cli_defaults_to_two_plus_two(self):
+        src = RUNNER.read_text(encoding="utf-8")
+        self.assertRegex(
+            src,
+            r'--jamba-transformer-layers", type=int, default=2',
+        )
+        self.assertRegex(
+            src,
+            r'--jamba-mamba-layers", type=int, default=2',
+        )