From c963abe9cd329d203a04d1236dc37ae8af127fab Mon Sep 17 00:00:00 2001 From: David Kelley Date: Sun, 4 Oct 2026 15:31:25 -0700 Subject: [PATCH 1/3] Store SeqNNMod.output_slice as an index tensor buffer A pandas Index output_slice (e.g. targets_df.index from hound_snp) made torch.compile skip SeqNNMod.forward entirely (graph break inside the head loop). As a non-persistent long-tensor buffer it traces into a single graph, follows .to(device), and leaves the state_dict unchanged. Co-Authored-By: Claude Opus 5.5 --- src/baskerville/seqnn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/baskerville/seqnn.py b/src/baskerville/seqnn.py index f151e7f..bb47ee9 100644 --- a/src/baskerville/seqnn.py +++ b/src/baskerville/seqnn.py @@ -805,7 +805,10 @@ def __init__( self.heads_gene_def = heads_gene_def self.output_stride = 1 self.output_crop_bp = 0 - self.output_slice = output_slice + # index tensor (not e.g. a pandas Index) so torch.compile can trace the slice + if output_slice is not None: + output_slice = torch.as_tensor(np.asarray(output_slice), dtype=torch.long) + self.register_buffer("output_slice", output_slice, persistent=False) global_vars = global_vars or {} self.global_vars = global_vars self.seq_length = seq_length From f7f7f514fa27820d822125e8a69e21617d66847e Mon Sep 17 00:00:00 2001 From: David Kelley Date: Sun, 4 Oct 2026 16:45:43 -0700 Subject: [PATCH 2/3] Make --mix_dtype take effect in SNP and ISM scoring, and warn on it - hound_snp never set SeqNN.mix_dtype, so SeqNN.__call__'s inner float32 autocast disabled the requested bf16/fp16: -m was a no-op. - SeqNN.__call__ now returns float32 predictions; reduced-precision outputs crashed scoring (numpy has no bfloat16) and summed imprecisely. - hound_snp_folds gains -m and passes it through. - -m is restricted to float32/bfloat16/float16 in all scripts. - snps.parse_mix_dtype warns that reduced precision can dominate variant effect scores (small alt - ref differences). Co-Authored-By: Claude Opus 5.5 --- src/baskerville/scripts/hound_ism_bed.py | 8 +-- .../scripts/hound_ism_bed_folds.py | 1 + src/baskerville/scripts/hound_ism_snp.py | 8 +-- .../scripts/hound_ism_snp_folds.py | 1 + src/baskerville/scripts/hound_snp.py | 11 +--- src/baskerville/scripts/hound_snp_folds.py | 11 ++++ src/baskerville/seqnn.py | 28 +++++----- src/baskerville/snps.py | 22 ++++++++ src/tests/test_ism_snp.py | 29 ++++++++++ src/tests/test_snp.py | 55 ++++++++++++++++++- 10 files changed, 140 insertions(+), 34 deletions(-) diff --git a/src/baskerville/scripts/hound_ism_bed.py b/src/baskerville/scripts/hound_ism_bed.py index 9e3db62..97ab31c 100755 --- a/src/baskerville/scripts/hound_ism_bed.py +++ b/src/baskerville/scripts/hound_ism_bed.py @@ -88,6 +88,7 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) parser.add_argument( @@ -151,12 +152,7 @@ def main(): args = parser.parse_args() # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = snps.parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = [snp_stat for snp_stat in args.snp_stats.split(",")] diff --git a/src/baskerville/scripts/hound_ism_bed_folds.py b/src/baskerville/scripts/hound_ism_bed_folds.py index 19b20a3..8546ae1 100755 --- a/src/baskerville/scripts/hound_ism_bed_folds.py +++ b/src/baskerville/scripts/hound_ism_bed_folds.py @@ -102,6 +102,7 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) ism_group.add_argument( diff --git a/src/baskerville/scripts/hound_ism_snp.py b/src/baskerville/scripts/hound_ism_snp.py index f8b1982..c15a7c4 100755 --- a/src/baskerville/scripts/hound_ism_snp.py +++ b/src/baskerville/scripts/hound_ism_snp.py @@ -87,6 +87,7 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) parser.add_argument( @@ -138,12 +139,7 @@ def main(): args = parser.parse_args() # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = snps.parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = args.snp_stats.split(",") diff --git a/src/baskerville/scripts/hound_ism_snp_folds.py b/src/baskerville/scripts/hound_ism_snp_folds.py index bbd058d..89271ae 100755 --- a/src/baskerville/scripts/hound_ism_snp_folds.py +++ b/src/baskerville/scripts/hound_ism_snp_folds.py @@ -99,6 +99,7 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) ism_group.add_argument( diff --git a/src/baskerville/scripts/hound_snp.py b/src/baskerville/scripts/hound_snp.py index fa6bf8f..c42d54b 100755 --- a/src/baskerville/scripts/hound_snp.py +++ b/src/baskerville/scripts/hound_snp.py @@ -17,9 +17,8 @@ import os import tempfile import shutil -import torch -from baskerville.snps import score_snps, score_gene_snps +from baskerville.snps import parse_mix_dtype, score_snps, score_gene_snps from baskerville.helpers.gcs_utils import ( upload_folder_gcs, download_rename_inputs, @@ -104,6 +103,7 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) parser.add_argument( @@ -207,12 +207,7 @@ def main(): args.targets_file = download_rename_inputs(args.targets_file, temp_dir) # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = args.snp_stats.split(",") diff --git a/src/baskerville/scripts/hound_snp_folds.py b/src/baskerville/scripts/hound_snp_folds.py index e8a4039..ccc45fc 100755 --- a/src/baskerville/scripts/hound_snp_folds.py +++ b/src/baskerville/scripts/hound_snp_folds.py @@ -108,6 +108,14 @@ def main(): type=int, help="Local window size in bp for targets with window=local", ) + snp_group.add_argument( + "-m", + "--mix_dtype", + dest="mix_dtype", + default="float32", + choices=["float32", "bfloat16", "float16"], + help="Mixed precision dtype", + ) snp_group.add_argument( "-n", "--norm", @@ -650,6 +658,9 @@ def build_snp_cmd(args, model_file, start_i, end_i, out_dir=None, fold_cross=Non if hasattr(args, "local_window") and args.local_window != 2048: cmd_parts.extend(["--local_window", str(args.local_window)]) + if hasattr(args, "mix_dtype") and args.mix_dtype != "float32": + cmd_parts.extend(["-m", args.mix_dtype]) + # Add normalization file if specified if hasattr(args, "norm_subdir") and args.norm_subdir and fold_cross is not None: norm_file = f"{args.models_dir}/{fold_cross}/{args.norm_subdir}/scores.h5" diff --git a/src/baskerville/seqnn.py b/src/baskerville/seqnn.py index bb47ee9..e1f97dd 100644 --- a/src/baskerville/seqnn.py +++ b/src/baskerville/seqnn.py @@ -15,6 +15,18 @@ from baskerville.types import BatchData, ModelOutput +def _float_output(out: ModelOutput, keep_gradients: bool) -> ModelOutput: + """Cast predictions to float32 (autocast may leave them reduced), detaching unless kept.""" + + def cast(x): + if x is None: + return None + x = x.float() + return x if keep_gradients else x.detach() + + return ModelOutput(coverage=cast(out.coverage), gene=cast(out.gene)) + + class SeqNN: """Sequence neural network model. @@ -142,12 +154,7 @@ def __call__( gene_out_mask=gene_out_mask, gene_presence=gene_presence, ) - if not keep_gradients: - out = ModelOutput( - coverage=out.coverage.detach() if out.has_coverage else None, - gene=out.gene.detach() if out.has_gene else None, - ) - outputs.append(out) + outputs.append(_float_output(out, keep_gradients)) # reverse complement if self.ensemble_rc: @@ -166,22 +173,17 @@ def __call__( gene_out_mask=rc_gene_out_mask, gene_presence=gene_presence, ) + out_rc = _float_output(out_rc, keep_gradients) # Process coverage: flip and apply strand pairing coverage_rc = None if out_rc.has_coverage: coverage_rc = torch.flip(out_rc.coverage, [2]) - if not keep_gradients: - coverage_rc = coverage_rc.detach() if strand_pair_hi is not None: coverage_rc = coverage_rc[:, strand_pair_hi, :] # Gene predictions: slices were transformed above, no post-hoc flipping needed - gene_rc = None - if out_rc.has_gene: - gene_rc = out_rc.gene - if not keep_gradients: - gene_rc = gene_rc.detach() + gene_rc = out_rc.gene outputs.append(ModelOutput(coverage=coverage_rc, gene=gene_rc)) diff --git a/src/baskerville/snps.py b/src/baskerville/snps.py index a96da3a..52968cf 100644 --- a/src/baskerville/snps.py +++ b/src/baskerville/snps.py @@ -1,6 +1,8 @@ import concurrent import json import os +import sys + import h5py import numpy as np import pandas as pd @@ -17,6 +19,24 @@ GENE_ASSAYS = {"RNA", "RNA3", "CAGE"} +MIX_DTYPES = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, +} + + +def parse_mix_dtype(name: str) -> torch.dtype: + """Map a --mix_dtype name to a torch dtype, warning on reduced precision.""" + if name != "float32": + print( + f"WARNING: --mix_dtype {name} adds rounding noise that can dominate " + "variant effect scores (small alt - ref differences). Use float32 " + "unless validated for your model and stats.", + file=sys.stderr, + ) + return MIX_DTYPES[name] + def gene_track_mask(targets_df) -> np.ndarray: """Boolean mask over targets selecting gene-appropriate tracks. @@ -164,6 +184,7 @@ def score_snps(args): seqnn_model = seqnn.SeqNN(params_model, output_slice=targets_df.index) seqnn_model.restore(args.model_file) seqnn_model.ensemble_rc = args.rc + seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) @@ -608,6 +629,7 @@ def score_gene_snps(args): seqnn_model = seqnn.SeqNN(params_model, output_slice=targets_df.index) seqnn_model.restore(args.model_file) seqnn_model.ensemble_rc = args.rc + seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) diff --git a/src/tests/test_ism_snp.py b/src/tests/test_ism_snp.py index de318c1..17654dc 100644 --- a/src/tests/test_ism_snp.py +++ b/src/tests/test_ism_snp.py @@ -299,3 +299,32 @@ def test_ism_snp_gene_track_selection(static_model_dir, stat): continue out_df = pd.read_csv(f"{temp_dir}/{name}", sep="\t", index_col=0) assert list(out_df.identifier) == expected_targets + + +def test_ism_snp_bfloat16(static_model_dir): + """-m bfloat16 runs (predictions cast to float32 for scoring) and warns.""" + test_dir = str(pathlib.Path(__file__).parent) + with tempfile.TemporaryDirectory() as temp_dir: + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_ism_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + temp_dir, + "-l", + "4", + "--stats", + "logSUM", + "-m", + "bfloat16", + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + f"{test_dir}/data/sc3_snps.vcf", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_ism_snp failed: {result.stderr}" + assert "WARNING: --mix_dtype bfloat16" in result.stderr diff --git a/src/tests/test_snp.py b/src/tests/test_snp.py index 2dd5b95..5542cf5 100644 --- a/src/tests/test_snp.py +++ b/src/tests/test_snp.py @@ -5,11 +5,15 @@ import pytest import subprocess +import argparse + import h5py import numpy as np import pandas as pd +import torch -from baskerville.snps import partition_snp_stats +from baskerville.scripts.hound_snp_folds import build_snp_cmd +from baskerville.snps import parse_mix_dtype, partition_snp_stats @pytest.fixture @@ -714,3 +718,52 @@ def test_partition_snp_stats(stats, expected): original = stats.copy() assert partition_snp_stats(stats) == expected assert stats == original + + +def test_parse_mix_dtype_warns_below_float32(capsys): + assert parse_mix_dtype("float32") is torch.float32 + assert "WARNING" not in capsys.readouterr().err + assert parse_mix_dtype("bfloat16") is torch.bfloat16 + assert "WARNING" in capsys.readouterr().err + + +def test_snp_folds_passes_mix_dtype(): + args = argparse.Namespace( + mix_dtype="bfloat16", params_file="p.json", vcf_file="v.vcf" + ) + assert "-m bfloat16" in build_snp_cmd(args, "m.pth", 0, 1) + args.mix_dtype = "float32" + assert "-m" not in build_snp_cmd(args, "m.pth", 0, 1).split() + + +def test_snp_mix_dtype_takes_effect(static_model_dir, test_vcf_file): + """-m bfloat16 must change the model's precision (it was once a silent no-op).""" + test_dir = str(pathlib.Path(__file__).parent) + scores = {} + with tempfile.TemporaryDirectory() as temp_dir: + for mix_dtype in ["float32", "bfloat16"]: + out_dir = f"{temp_dir}/{mix_dtype}" + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + out_dir, + "--stats", + "logSUM", + "-m", + mix_dtype, + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + test_vcf_file, + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_snp failed: {result.stderr}" + assert ("WARNING: --mix_dtype" in result.stderr) == (mix_dtype != "float32") + with h5py.File(f"{out_dir}/scores.h5", "r") as h5: + scores[mix_dtype] = h5["cov/logSUM"][:] + assert not np.array_equal(scores["float32"], scores["bfloat16"]) From 99a7f9f1a26c5d24e8ec9e7f02087d845b090975 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Sun, 4 Oct 2026 18:07:59 -0700 Subject: [PATCH 3/3] Add --compile to SNP and ISM scoring scripts hound_snp, hound_ism_snp, hound_ism_bed and their _folds wrappers gain --compile, which calls SeqNN.compile() after restore. Default mode only: reduce-overhead's CUDA graphs overwrite outputs held across calls (ref predictions). SeqNN.compile only queries CUDA capability on CUDA devices. Co-Authored-By: Claude Opus 5.5 --- src/baskerville/scripts/hound_ism_bed.py | 8 ++++ .../scripts/hound_ism_bed_folds.py | 9 +++++ src/baskerville/scripts/hound_ism_snp.py | 8 ++++ .../scripts/hound_ism_snp_folds.py | 9 +++++ src/baskerville/scripts/hound_snp.py | 6 +++ src/baskerville/scripts/hound_snp_folds.py | 9 +++++ src/baskerville/seqnn.py | 8 ++-- src/baskerville/snps.py | 4 ++ src/tests/test_snp.py | 38 +++++++++++++++++++ 9 files changed, 96 insertions(+), 3 deletions(-) diff --git a/src/baskerville/scripts/hound_ism_bed.py b/src/baskerville/scripts/hound_ism_bed.py index 97ab31c..3116c42 100755 --- a/src/baskerville/scripts/hound_ism_bed.py +++ b/src/baskerville/scripts/hound_ism_bed.py @@ -91,6 +91,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-o", "--out_dir", @@ -215,6 +221,8 @@ def main(): seqnn_model.ensemble_shifts = args.shifts seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() + if args.compile: + seqnn_model.compile() output_length = seqnn_model.output_length() output_stride = seqnn_model.output_stride() diff --git a/src/baskerville/scripts/hound_ism_bed_folds.py b/src/baskerville/scripts/hound_ism_bed_folds.py index 8546ae1..43caab0 100755 --- a/src/baskerville/scripts/hound_ism_bed_folds.py +++ b/src/baskerville/scripts/hound_ism_bed_folds.py @@ -105,6 +105,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + ism_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) ism_group.add_argument( "-o", dest="out_dir", @@ -541,6 +547,9 @@ def build_ism_bed_cmd(args, model_file, out_dir): if args.mix_dtype != "float32": cmd_parts.extend(["-m", args.mix_dtype]) + if args.compile: + cmd_parts.append("--compile") + cmd_parts.extend(["-o", out_dir]) if args.rc: diff --git a/src/baskerville/scripts/hound_ism_snp.py b/src/baskerville/scripts/hound_ism_snp.py index c15a7c4..08d49c3 100755 --- a/src/baskerville/scripts/hound_ism_snp.py +++ b/src/baskerville/scripts/hound_ism_snp.py @@ -90,6 +90,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-o", dest="out_dir", @@ -199,6 +205,8 @@ def main(): seqnn_model.ensemble_shifts = args.shifts seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() + if args.compile: + seqnn_model.compile() output_length = seqnn_model.output_length() output_stride = seqnn_model.output_stride() diff --git a/src/baskerville/scripts/hound_ism_snp_folds.py b/src/baskerville/scripts/hound_ism_snp_folds.py index 89271ae..972a2ae 100755 --- a/src/baskerville/scripts/hound_ism_snp_folds.py +++ b/src/baskerville/scripts/hound_ism_snp_folds.py @@ -102,6 +102,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + ism_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) ism_group.add_argument( "-o", dest="out_dir", @@ -539,6 +545,9 @@ def build_ism_snp_cmd(args, model_file, out_dir): if args.mix_dtype != "float32": cmd_parts.extend(["-m", args.mix_dtype]) + if args.compile: + cmd_parts.append("--compile") + cmd_parts.extend(["-o", out_dir]) if args.rc: diff --git a/src/baskerville/scripts/hound_snp.py b/src/baskerville/scripts/hound_snp.py index c42d54b..907c450 100755 --- a/src/baskerville/scripts/hound_snp.py +++ b/src/baskerville/scripts/hound_snp.py @@ -106,6 +106,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-n", "--norm", diff --git a/src/baskerville/scripts/hound_snp_folds.py b/src/baskerville/scripts/hound_snp_folds.py index ccc45fc..80dcfca 100755 --- a/src/baskerville/scripts/hound_snp_folds.py +++ b/src/baskerville/scripts/hound_snp_folds.py @@ -116,6 +116,12 @@ def main(): choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + snp_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) snp_group.add_argument( "-n", "--norm", @@ -661,6 +667,9 @@ def build_snp_cmd(args, model_file, start_i, end_i, out_dir=None, fold_cross=Non if hasattr(args, "mix_dtype") and args.mix_dtype != "float32": cmd_parts.extend(["-m", args.mix_dtype]) + if hasattr(args, "compile") and args.compile: + cmd_parts.append("--compile") + # Add normalization file if specified if hasattr(args, "norm_subdir") and args.norm_subdir and fold_cross is not None: norm_file = f"{args.models_dir}/{fold_cross}/{args.norm_subdir}/scores.h5" diff --git a/src/baskerville/seqnn.py b/src/baskerville/seqnn.py index e1f97dd..bbbc194 100644 --- a/src/baskerville/seqnn.py +++ b/src/baskerville/seqnn.py @@ -202,10 +202,12 @@ def __call__( return ModelOutput(coverage=coverage_avg, gene=gene_avg) def compile(self): - """Compile model for faster inference. - (Not working well with augmentation ensembling.) + """Compile model for faster inference; call after restore. + + Default mode: reduce-overhead's CUDA graphs overwrite outputs that + callers hold across calls (e.g. ref predictions). """ - if torch.cuda.get_device_capability()[0] < 7: + if self.device == "cuda" and torch.cuda.get_device_capability()[0] < 7: print("Warning: CUDA device capability < 7.0, skipping compilation.") else: self.model = torch.compile(self.model) diff --git a/src/baskerville/snps.py b/src/baskerville/snps.py index 52968cf..635f021 100644 --- a/src/baskerville/snps.py +++ b/src/baskerville/snps.py @@ -187,6 +187,8 @@ def score_snps(args): seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) + if args.compile: + seqnn_model.compile() # shift outside seqnn num_shifts = len(args.shifts) @@ -632,6 +634,8 @@ def score_gene_snps(args): seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) + if args.compile: + seqnn_model.compile() # shift outside seqnn num_shifts = len(args.shifts) diff --git a/src/tests/test_snp.py b/src/tests/test_snp.py index 5542cf5..f54c303 100644 --- a/src/tests/test_snp.py +++ b/src/tests/test_snp.py @@ -736,6 +736,11 @@ def test_snp_folds_passes_mix_dtype(): assert "-m" not in build_snp_cmd(args, "m.pth", 0, 1).split() +def test_snp_folds_passes_compile(): + args = argparse.Namespace(compile=True, params_file="p.json", vcf_file="v.vcf") + assert "--compile" in build_snp_cmd(args, "m.pth", 0, 1).split() + + def test_snp_mix_dtype_takes_effect(static_model_dir, test_vcf_file): """-m bfloat16 must change the model's precision (it was once a silent no-op).""" test_dir = str(pathlib.Path(__file__).parent) @@ -767,3 +772,36 @@ def test_snp_mix_dtype_takes_effect(static_model_dir, test_vcf_file): with h5py.File(f"{out_dir}/scores.h5", "r") as h5: scores[mix_dtype] = h5["cov/logSUM"][:] assert not np.array_equal(scores["float32"], scores["bfloat16"]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="compile targets GPU") +def test_snp_compile_matches_eager(static_model_dir, test_vcf_file): + """--compile runs and reproduces eager scores up to fusion reassociation.""" + test_dir = str(pathlib.Path(__file__).parent) + scores = {} + with tempfile.TemporaryDirectory() as temp_dir: + for compile_flag in [[], ["--compile"]]: + out_dir = f"{temp_dir}/{len(compile_flag)}" + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + out_dir, + "--stats", + "logSUM", + "--rc", + *compile_flag, + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + test_vcf_file, + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_snp failed: {result.stderr}" + with h5py.File(f"{out_dir}/scores.h5", "r") as h5: + scores[len(compile_flag)] = h5["cov/logSUM"][:].astype("float32") + np.testing.assert_allclose(scores[1], scores[0], rtol=1e-2, atol=1e-2)