Make --mix_dtype work in SNP/ISM scoring; add --compile

src/baskerville/scripts/hound_ism_bed.py

@@ -88,8 +88,15 @@ def main():
         "--mix_dtype",
         dest="mix_dtype",
         default="float32",
+        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",
@@ -151,12 +158,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(",")]
@@ -219,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()

src/baskerville/scripts/hound_ism_bed_folds.py

@@ -102,8 +102,15 @@ def main():
         "--mix_dtype",
         dest="mix_dtype",
         default="float32",
+        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",
@@ -540,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:

src/baskerville/scripts/hound_ism_snp.py

@@ -87,8 +87,15 @@ def main():
         "--mix_dtype",
         dest="mix_dtype",
         default="float32",
+        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",
@@ -138,12 +145,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(",")
 
@@ -203,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()

src/baskerville/scripts/hound_ism_snp_folds.py

@@ -99,8 +99,15 @@ def main():
         "--mix_dtype",
         dest="mix_dtype",
         default="float32",
+        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",
@@ -538,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:

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,8 +103,15 @@ def main():
         "--mix_dtype",
         dest="mix_dtype",
         default="float32",
+        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",
@@ -207,12 +213,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(",")
 

src/baskerville/scripts/hound_snp_folds.py

@@ -108,6 +108,20 @@ 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(
+        "--compile",
+        default=False,
+        action="store_true",
+        help="Compile the model with torch.compile",
+    )
     snp_group.add_argument(
         "-n",
         "--norm",
@@ -650,6 +664,12 @@ 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])
+
+    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"

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))
 
@@ -200,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)
@@ -805,7 +809,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

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,8 +184,11 @@ 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)
+    if args.compile:
+        seqnn_model.compile()
 
     # shift outside seqnn
     num_shifts = len(args.shifts)
@@ -608,8 +631,11 @@ 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)
+    if args.compile:
+        seqnn_model.compile()
 
     # shift outside seqnn
     num_shifts = len(args.shifts)

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