alphagenome_pt: assay_ids=None crashed, and the splice head was being read as logits

[lucapinello]

Fixes the defect recorded as a known issue in #238, plus a second one that fixing the first exposed, plus
a third that only showed up end-to-end.

The crash

predict_variant_effect(..., assay_ids=None) raised invalid literal for int() with base 10: 'logits'
on this backend, while an explicit list of ATAC ids worked. The default is all 5,168 tracks and the 4
SPLICE_SITES tracks are always among them, so every default-argument call failed.

Cause: head-output dicts are keyed by two unrelated things. Most heads key by resolution
({1: ..., 128: ...}); SpliceSitesClassificationHead keys by tensor kind
({"logits": ..., "probs": ...}). The extraction loop reassigned its own res variable to whichever key
it picked, then called int(res). Resolution and tensor selection are now separate questions, asked
separately, in _select_head_tensor and _as_resolution.

The quieter bug underneath it

The old fallback was next(iter(head_out.keys())) — dict insertion order, which for
{"logits", "probs"} is logits.

The JAX reference returns {'logits': logits, 'predictions': softmax(logits)} and treats 'predictions'
as the prediction: alphagenome_research/model/model.py:281 reads exactly that key to derive splice
junctions. The pt port names the same tensor probs. So the two backends would have disagreed on those
4 tracks — unbounded logits against probabilities in [0, 1] — had the crash not masked it. Selection now
prefers probs explicitly and no longer depends on insertion order.

The third one: my first fix was incomplete, and the unit tests did not notice

After the first commit, all 13 tests passed and assay_ids=None still crashed end-to-end.

use_environment=True — this backend's default — does not call _predict_raw. It executes
alphagenome_pt_source/templates/predict_template.py in the per-oracle conda env, and that template
carried its own copy of the same six lines. I had fixed the path this backend does not use by default.

The template now imports the two helpers instead of restating them. Two copies of one rule is the failure
mode this repo keeps rediscovering — the Sei normalizer and the Jupyter kernel were both correct code
nothing called, and the Sei background builder nearly shipped a second copy of the nucleosome correction.

Tests

16 tests, and 9 of them fail against the pre-fix code (verified by reverting each part):

• 6 on the oracle logic — the crash half via _as_resolution, the correctness half via probs-vs-logits
  and an insertion-order test that pins the answer against both dict orderings.
• 3 on the template — it must import the helpers, must not contain int(res), must not reimplement the
  next(iter(head_out.keys())) fallback.

The oracle tests need no torch or GPU because the tensors are opaque to the selection logic. That
isolation is deliberate but not sufficient on its own, and this PR is the proof: the guard written for
Sei's equivalent assay_ids=None defect asserted at the validation layer and did not catch this one, and
my own unit tests here passed while the default path stayed broken. Hence the source-level template
guards and the end-to-end run below.

Fast suite 2,139 passed / 35 skipped.

End-to-end verification

predict_variant_effect in env mode over 9 ids spanning all 9 heads, including SPLICE_SITES/donor/+:
9 tracks returned, no crash. The 128-bp heads still return 8,192 bins against 1,048,576 for the 1-bp
heads, so _as_resolution reports real resolutions rather than defaulting.

Correction to an earlier version of this description

An earlier revision claimed heads= filtering was broken in the in-process path, on the strength of a
predict() call dying with Expected size for first two dimensions of batch2 tensor to be: [8, 16] but got: [8, 1]. That claim was wrong and is withdrawn. I then blamed interval length, which was also
wrong — it reproduced at 1 kb and at full window, in-process and in env mode.

Instrumenting the template to log what actually reached the model gave the answer:

len(sequence)=24        heads=('atac',)        <- predict()
len(sequence)=1048576   heads=(all 9)          <- predict_variant_effect()

len("chr1:109274468-109275468") == 24. predict() takes a str as literal DNA, so my coordinate
string was one-hot encoded as 24 non-ACGT bases. My probe misused the API; the heads filter and the
interval were never involved. All four coordinate-taking oracles behave identically here, so it is
intended and uniform.

The one residual, pre-existing trap worth a follow-up: the same string is a valid region for
predict_variant_effect and literal DNA for predict(), and the mistake surfaces as a torch shape
error rather than "this does not look like DNA". Out of scope here — it is uniform across oracles and
predates this change — but it cost me real time, which is the argument for validating it.

This also closes #238's last open row

With the crash fixed, alphagenome_pt cross-process determinism is finally measurable, and it is
bit-exact: 2/2 pairs, max|diff| = 0.000e+00 over 9 tracks spanning all 9 heads (288 values).

So all nine oracles are now measured cross-process and Enformer is the only one that drifts. That is
the strongest form of the argument for accepting the drift rather than rebuilding the nulls: ChromBPNet,
the other TF-backed oracle, is bit-exact, so the risk is one oracle wide and near_ties_at_cutoff
already covers the one way it reaches a reader.

Worth noting how the gap read: "not measured" looked like a scheduling gap in a table when it was
actually a live bug in the default code path. The blank was the evidence.