fix(openfold3): match AF3 pTM/ipTM reduction and confidence bin centers

[ashanehsazzadeh-dev]

Description

Two deviations in bionemo_ir/pipeline/models/openfold3/postprocessor.py (line numbers at bbf5f9ec) from the AlphaFold 3 definitions (Abramson et al., Nature 2024, SI §5.9.1, eqs. 17–18) and from upstream OpenFold3 (openfold3/core/metrics/confidence.py):

1. pTM / ipTM reduction. _tm_score_from_pae_logits (lines 262–282) returns the mean of the expected TM term over all token pairs (pTM) or over all inter-chain pairs (ipTM). AF3 and upstream compute_ptm take, for each aligned token i, the mean over scored tokens j and then the maximum over i (restricted to tokens with a valid frame). The pair mean is a lower bound of that quantity, so 0.1.0 reports systematically low pTM/ipTM — not comparable with the usual ipTM thresholds or with 0.8·ipTM + 0.2·pTM — while the PAE matrix itself is correct. The Boltz-2, Protenix and OpenFold2 paths in this repository already use row-mean-then-max with bin midpoints; the OpenFold3 post-processor was the exception.
2. Bin centers. The TM term (line 270), _compute_pae (298), _compute_plddt (209) and _select_best_sample (188) weight bin probabilities with torch.linspace(lo, hi, n_bins) end points instead of bin midpoints (upstream get_bin_centers: 0.25 … 31.75 Å for PAE, 0.01 … 0.99 for pLDDT). For PAE and pLDDT this is an exact increasing affine map, so the pipeline's own best-sample selection (mean-pLDDT argmax) is unchanged by this PR.

Changes (one PR, two commits, bisectable)

1. fix(openfold3): match AF3 pTM/ipTM reduction (row mean, max over aligned tokens) — reduction fix, TM term at bin midpoints, optional has_frame argument. At the call site the mask is derived from the existing is_atomized feature (polymer tokens are frame-eligible; ligand atoms / ions are scored as j but not aligned on). This equals upstream's mask for polymer residues and single-atom ligands and is a lower bound for multi-atom ligands (upstream's get_token_frame_atoms is not ported here). If no token is eligible the score is nan (→ None in get_scores()), like the existing single-chain ipTM case. get_scores() gains one optional key, ptm_frame_mask, recording which convention produced the values — happy to drop it if you prefer the scores schema unchanged.
2. fix(openfold3): use bin midpoints for PAE and pLDDT expectations — _compute_pae, _compute_plddt, _select_best_sample. Shifts every PAE entry by ≤ 0.25 Å and every pLDDT / B-factor / ModelCIF QA value by ≤ 1; kept separate from commit 1 because it touches many outputs but changes no ranking.

Conversions for existing 0.1.0 outputs: PAE_new = 63/64 · PAE_0.1.0 + 0.25 Å; pLDDT_new = 49/50 · pLDDT_0.1.0 + 1 (0–100 scale); max_pae ceiling 32.0 → 31.75 Å; pTM / ipTM have no closed form — recompute from the PAE logits.

Example (public PDB entries)

Same PAE logits through the 0.1.0 and the patched post-processing, next to upstream OpenFold3 run separately on the same inputs (one seed, mean over 5 diffusion samples, so the upstream column agrees within sample-to-sample scatter rather than exactly):

entry	tokens	pTM 0.1.0 → patched	pTM upstream	ipTM 0.1.0 → patched	ipTM upstream
9H1J (two protein chains)	172	0.552 → 0.745	0.744	0.307 → 0.415	0.410
9RVG (homodimer + 2 citrate)	410	0.636 → 0.831	0.834	0.586 → 0.812	0.818

On both entries pLDDT and PAE change only by the affine maps above, and the pipeline's selected sample is the same before and after.

Tests

tests/pipeline/models/openfold3/test_openfold3_ptm.py (new; synthetic inputs): max-over-aligned-tokens for pTM (0.856 vs 0.037 with the pair mean) and ipTM (inter-chain rows only); TM term at bin midpoints with the AF3 d0 (incl. the N < 19 clip); has_frame semantics incl. the all-masked / ligand-only case; (B, S, N, N, 64) slicing; PAE = 0.25 + 0.5k Å and pLDDT = 2k + 1 for one-hot bins; best-sample argmax invariance; parity with upstream compute_ptm, probs_to_expected_error and compute_plddt (|Δ| < 1e-5; skipped when the 3rdparty/openfold-3 submodule is not checked out). With both commits: 15 pass, 1 skipped (an env-gated report helper); on current main 13 of these fail. ruff check and ruff format --check are clean. The tests were run on CPU here — please let CI confirm on GPU.

Type of change

• Bug fix (non-breaking)
• New feature (non-breaking)
• Breaking change
• Refactor
• Documentation
• Build / CI

Checklist

• My commits are signed off (DCO): git commit -s
• I have read the Contributing Guidelines
• For a new feature or breaking change, an issue was filed and approved first — n/a (bug fix)
• I added or updated tests, and they pass locally
• I updated documentation as needed — n/a

[ducta3141]

cc @thanhnamitit

[ducta3141]

This return is redundant, we should leave to the original.

[ducta3141]

Re-check with OSS and AF3 and confirm this is the real bug for computing scores.

[ducta3141]

This below is recommending fixes with efficient memory, basically we should reuse atomize_utils.py::get_token_frame_mask under the model-level:

diff --git a/bionemo_ir/_torch/modules/openfold3/utils/atomize_utils.py b/bionemo_ir/_torch/modules/openfold3/utils/atomize_utils.py
index 2a11015..529a919 100644
--- a/bionemo_ir/_torch/modules/openfold3/utils/atomize_utils.py
+++ b/bionemo_ir/_torch/modules/openfold3/utils/atomize_utils.py
@@ -295,21 +295,87 @@ def get_token_atom_index_offset(atom_name: str, restype: torch.Tensor):
     return token_atom_index_offset, token_atom_mask
 
 
-def get_token_frame_atoms(
+def _closest_atoms_to_start_atoms(
+    x: torch.Tensor,
+    atom_mask: torch.Tensor,
+    atom_asym_id: torch.Tensor,
+    start_atom_index: torch.Tensor,
+    eps: float,
+    inf: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Indices of the two closest same-chain atoms to each token's start atom.
+
+    A dense neighbour search over all atom pairs computes ``N_atom`` rows of which
+    only the ``N_token`` start-atom rows are ever read, so the query axis here is the
+    tokens. Distances and the pair mask are formed exactly as the dense search forms
+    them, so the selected indices are unchanged, but the working set is
+    ``N_token x N_atom`` per diffusion sample rather than ``N_atom^2`` -- the
+    difference between single-digit GB and tens of GB on a large complex.
+
+    Args:
+        x:
+            [*, N_atom, 3] Atom positions
+        atom_mask:
+            [*, N_atom] Atom mask
+        atom_asym_id:
+            [*, N_atom] Chain index, broadcast to atoms
+        start_atom_index:
+            [*, N_token] Index of the first atom of each token
+        eps:
+            Small constant for numerical stability
+        inf:
+            Large constant for numerical stability
+
+    Returns:
+        ([*, N_token], [*, N_token])
+        Indices of the closest and second closest atom to each token's start atom
+    """
+    leading_shape = x.shape[:-2]
+    atom_mask = torch.broadcast_to(atom_mask, (*leading_shape, atom_mask.shape[-1]))
+    atom_asym_id = torch.broadcast_to(atom_asym_id, (*leading_shape, atom_asym_id.shape[-1]))
+
+    # Position, mask and chain of the query (start) atoms
+    start_x = torch.gather(x, dim=-2, index=start_atom_index.unsqueeze(-1).expand(*start_atom_index.shape, 3))
+    start_atom_mask = torch.gather(atom_mask, dim=-1, index=start_atom_index)
+    start_asym_id = torch.gather(atom_asym_id, dim=-1, index=start_atom_index)
+
+    # Pairwise mask over (start atom, atom): both present, and within the same chain
+    # [*, N_token, N_atom]
+    pair_mask = start_atom_mask[..., None] * atom_mask[..., None, :]
+    pair_mask = pair_mask * (start_asym_id[..., None] == atom_asym_id[..., None, :])
+
+    # Distance from every start atom to every atom
+    # [*, N_token, N_atom]
+    d = torch.sum(eps + (start_x[..., None, :] - x[..., None, :, :]) ** 2, dim=-1) ** 0.5
+    d = d * pair_mask + inf * (1 - pair_mask)
+
+    # Index 0 is the start atom itself, so 1 and 2 are its two closest neighbours
+    _, closest_atom_index = torch.topk(d, k=3, dim=-1, largest=False)
+    return closest_atom_index[..., 1], closest_atom_index[..., 2]
+
+
+def get_token_frame_mask(
     batch: dict,
     x: torch.Tensor,
     atom_mask: torch.Tensor,
     angle_threshold: float = 25.0,
     eps: float = 1e-8,
     inf: float = 1e9,
-):
+) -> torch.Tensor:
     """
-    Extract frame atoms per token, which returns
+    Mask of tokens whose frame is valid, from the frame atoms
         -   (N, Ca, C) for standard amino acid residues
         -   (C3', C1', C4') for standard nucleotide residues
         -   closest neighbors for atomized tokens (modified residues and ligands),
             subject to additional angle and chain constraints from Subsection 4.3.2
 
+    A frame is valid when its three atoms are present, lie in one chain, and -- for
+    atomized tokens, whose frame comes from nearest neighbours rather than a known
+    backbone -- span an angle within ``angle_threshold`` of neither 0 nor 180 degrees.
+    This is the ``has_frame`` input of the pTM / ipTM outer maximum (AF3 SI 5.9.1):
+    only a token with a frame can be the aligned token. The frame atom positions
+    themselves are used only to test that angle and are not returned.
+
     Args:
         batch:
             Feature dictionary
@@ -324,37 +390,28 @@ def get_token_frame_atoms(
         inf:
             Large constant for numerical stability
     Returns:
-        phi:
-            ([*, N_token, 3], [*, N_token, 3], [*, N_token, 3])
-            Tuple of three frame atoms
         valid_frame_mask:
             [*, N_token] Mask denoting valid frames
     """
-    # Create pairwise atom mask
-    pair_mask = atom_mask[..., None] * atom_mask[..., None, :]
-
-    # Update pairwise atom mask
-    # Restrict to atoms within the same chain
+    # Chain index per atom, to restrict frames to atoms within the same chain
     atom_asym_id = broadcast_token_feat_to_atoms(
         token_mask=batch["token_mask"],
         num_atoms_per_token=batch["num_atoms_per_token"],
         token_feat=batch["asym_id"],
     )
-    atom_asym_id_mask = atom_asym_id[..., None] == atom_asym_id[..., None, :]
-    pair_mask = pair_mask * atom_asym_id_mask
-
-    # Compute distance matrix
-    # [*, N_atom, N_atom]
-    d = torch.sum(eps + (x[..., None, :] - x[..., None, :, :]) ** 2, dim=-1) ** 0.5
-    d = d * pair_mask + inf * (1 - pair_mask)
 
     # Find indices of two closest atoms for start atoms
     # [*, N_token]
     start_atom_index = batch["start_atom_index"].long()
     start_atom_index = start_atom_index.expand(*x.shape[:-2], start_atom_index.shape[-1])
-    _, closest_atom_index = torch.topk(d, k=3, dim=-1, largest=False)
-    a_index = torch.gather(closest_atom_index[..., 1], dim=-1, index=start_atom_index)
-    c_index = torch.gather(closest_atom_index[..., 2], dim=-1, index=start_atom_index)
+    a_index, c_index = _closest_atoms_to_start_atoms(
+        x=x,
+        atom_mask=atom_mask,
+        atom_asym_id=atom_asym_id,
+        start_atom_index=start_atom_index,
+        eps=eps,
+        inf=inf,
+    )
 
     # Construct indices of atoms used for frame construction
     # [*, N_token]
@@ -401,7 +458,8 @@ def get_token_frame_atoms(
         },
     }
 
-    # Extract coordinates
+    # Extract chain, presence and coordinates of each frame atom. The coordinates
+    # serve only the angle test below; they are not part of the result.
     for key in frame_atoms:
         frame_atoms[key].update(
             {
@@ -457,11 +515,4 @@ def get_token_frame_atoms(
     )
 
     # Compute final valid frame mask
-    valid_frame_mask = valid_frame_mask_angle * valid_frame_mask_atom * valid_frame_mask_asym_id
-    phi = (
-        frame_atoms["a"]["atom_positions"],
-        frame_atoms["b"]["atom_positions"],
-        frame_atoms["c"]["atom_positions"],
-    )
-
-    return phi, valid_frame_mask
+    return valid_frame_mask_angle * valid_frame_mask_atom * valid_frame_mask_asym_id

[ducta3141]

At the confidence.py, the computation of valid_frame_mask from get_token_frame_atoms:

diff --git a/bionemo_ir/_torch/modules/openfold3/confidence.py b/bionemo_ir/_torch/modules/openfold3/confidence.py
index 6b8c081..98e9f58 100644
--- a/bionemo_ir/_torch/modules/openfold3/confidence.py
+++ b/bionemo_ir/_torch/modules/openfold3/confidence.py
@@ -22,6 +22,7 @@ from bionemo_ir._torch.layers.linear import Linear
 from bionemo_ir._torch.layers.transformers.pairformer import PairformerModule
 from bionemo_ir._torch.modules.openfold3.utils.atomize_utils import (
     broadcast_token_feat_to_atoms,
+    get_token_frame_mask,
     get_token_representative_atoms,
     max_atom_per_token_masked_select,
 )
@@ -657,6 +658,9 @@ class AuxiliaryHeadsAllAtom(nn.Module):
                         Predicted binned PLDDT logits
                     "pae_logits" ([*, N_token, N_token, 64]):
                         Predicted binned PAE logits
+                    "valid_frame_mask" ([*, N_token]):
+                        Tokens with a valid frame, the ``has_frame`` input of the
+                        pTM / ipTM outer maximum. Present with "pae_logits" only.
                     "pde_logits" ([*, N_token, N_token, 64]):
                         Predicted binned PDE logits
                     "experimentally_resolved_logits" ([*, N_atom, 2]):
@@ -721,6 +725,12 @@ class AuxiliaryHeadsAllAtom(nn.Module):
 
         if self.config.pae.enabled:
             aux_out["pae_logits"] = self.pae(zij).to(device=out_device)
+            # has_frame for the pTM / ipTM outer maximum, from the sampled
+            # coordinates, so it is per diffusion sample. Only the PAE head feeds
+            # pTM / ipTM, so nothing needs it when that head is off.
+            aux_out["valid_frame_mask"] = get_token_frame_mask(
+                batch=batch, x=atom_positions_predicted, atom_mask=batch["atom_mask"]
+            ).to(device=out_device)
 
         aux_out["pde_logits"] = pde_logits.to(device=out_device)

[ducta3141]

And below suggestions on the file postprocess.py will move tensors computation for scores on the GPU devices.

[ducta3141]

Suggested change

	plddt_per_atom = _plddt_per_atom(output)
	raw_pos = torch.as_tensor(output["atom_positions_predicted"])
	if raw_pos.dim() == 4:
	# Multiple diffusion samples — select best by mean pLDDT
	best_idx = _select_best_sample(plddt_per_atom)
	best_pos = raw_pos[0, best_idx]
	elif raw_pos.dim() == 3:
	best_idx = 0
	best_pos = raw_pos[0]
	else:
	best_idx = 0
	best_pos = raw_pos
	# Copy only the chosen sample, not all of them
	best_pos = best_pos.cpu().numpy() # (N_atoms, 3)

[ducta3141]

The MR should be rebased first to the latest.

Suggested change

	# --- Atom-to-token mapping ---
	atom_to_token = _cpu(batch["atom_to_token_index"]).squeeze(0).numpy() # (N_atoms,)
	# --- Decode atom names ---
	raw_atom_names = batch.get("ref_atom_name_chars")
	if raw_atom_names is not None:
	raw_atom_names = _cpu(raw_atom_names).squeeze(0)
	flat_atom_names = decode_atom_name_chars(raw_atom_names, atom_mask_bool)
	# --- Remap into the shared atom layout ---
	# The universe covers protein backbone+sidechain, nucleic backbone,
	# nucleobases, and common ligand atom labels. Ligand tokens are
	# atomized (one atom = one token), so
	# writing each atom to its own (token, atom-name-slot) does not
	# collide with neighbouring tokens: there's no shared vocabulary
	# problem here because each ligand atom owns its own row.
	atom_positions, atom_mask_out = scatter_flat_atoms_to_folding_layout(
	best_pos,
	atom_to_token,
	flat_atom_names,
	atom_mask_bool,
	n_tokens,
	)

[ducta3141]

        pae_logits = _pae_logits(output, best_idx, n_tokens)
        has_frame = _frame_mask(output, best_idx, n_tokens)
        ptm = _compute_ptm(pae_logits, n_tokens, has_frame=has_frame)
        iptm = _compute_iptm(pae_logits, n_tokens, chain_indices, has_frame=has_frame)
        pae = _compute_pae(pae_logits)

[ducta3141]

Remove this part.

[ducta3141]

Suggested change

	def _aligned_token_mask(batch: dict[str, Any], n_tokens: int) -> torch.Tensor | None:
	"""Tokens eligible as the aligned token ``i`` in the pTM / ipTM max (``has_frame``).
	Interim stand-in for the coordinate-based frame validity of upstream OpenFold3
	(``openfold3/core/utils/atomize_utils.py::get_token_frame_atoms``): polymer
	tokens are eligible; atomized tokens (``batch["is_atomized"]``: ligand atoms,
	ions) are scored as ``j`` but never used as aligned tokens. Upstream additionally
	admits atomized tokens whose nearest-neighbour local frame is valid and requires
	the backbone frame atoms of polymer residues to be present. Returns ``None``
	(every token eligible) only when the flag is absent from the batch; an input
	without polymer tokens (ligand-only query) yields an all-False mask, for which
	pTM / ipTM are reported as NaN (see ``_tm_score_from_pae_logits``).
	"""
	is_atomized = batch.get("is_atomized")
	if is_atomized is None:
	return None
	return ~_cpu(is_atomized).reshape(-1)[:n_tokens].bool()
	def _frame_mask(output: dict, best_idx: int, n_tokens: int) -> torch.Tensor | None:
	"""``has_frame`` for the pTM / ipTM maximum, as emitted by the confidence head.
	The head computes it from the sampled coordinates, so it carries the diffusion
	sample axis and arrives as 0/1 floats in the output dtype. Absent when the PAE
	head is disabled, in which case there is no pTM to restrict. Stays on its
	original device, alongside the logits it will gate.
	"""
	mask = output.get("valid_frame_mask")
	if mask is None:
	return None
	mask = torch.as_tensor(mask)
	if mask.dim() == 3:
	mask = mask[0, best_idx]
	elif mask.dim() == 2:
	mask = mask[0]
	return mask[:n_tokens].bool()
	def _pae_logits(output: dict, best_idx: int, n_tokens: int) -> torch.Tensor | None:
	"""The selected sample's PAE logits, cropped to the real tokens.
	Left on its original device: ``torch.as_tensor`` preserves it, unlike
	``_cpu``. The engine calls the post-processor with the model's own output, so
	this is normally GPU memory, and every consumer reduces it.
	"""
	logits = output.get("pae_logits")
	if logits is None:
	return None
	logits = torch.as_tensor(logits)
	if logits.dim() == 5:
	logits = logits[0, best_idx] # (N_tokens, N_tokens, n_bins)
	elif logits.dim() == 4:
	logits = logits[0]
	return logits[:n_tokens, :n_tokens]
	def _plddt_per_atom(output: dict) -> torch.Tensor | None:
	"""Per-atom pLDDT on a 0-100 scale, as ``(n_samples, N_atom)``.
	Both the diffusion-sample choice and the reported per-token score are this
	same expectation, so it runs once here instead of once in each. Stays on the
	logits' device; only the selected row is ever copied to the host.
	"""
	logits = output.get("plddt_logits")
	if logits is None:
	return None
	logits = torch.as_tensor(logits)
	if logits.dim() == 4:
	logits = logits[0] # (B, S, N_atom, n_bins) -> (S, N_atom, n_bins)
	elif logits.dim() == 3:
	logits = logits[:1] # (B, N_atom, n_bins) -> a single sample
	if logits.dim() == 2:
	logits = logits.unsqueeze(0) # (N_atom, n_bins) -> a single sample
	probs = torch.softmax(logits.float(), dim=-1)
	bin_centers = _bin_centers(0.0, 1.0, probs.shape[-1]).to(device=probs.device)
	return (probs * bin_centers).sum(dim=-1) * 100.0

[ducta3141]

def _select_best_sample(plddt_per_atom: torch.Tensor | None) -> int:
    """Select best diffusion sample by mean pLDDT.

    Reduces on the device and brings back only the winning index.
    """
    if plddt_per_atom is None:
        return 0
    return int(plddt_per_atom.mean(dim=-1).argmax())

[ducta3141]

Suggested change

	plddt_per_atom = (probs * bin_centers).sum(dim=-1).numpy() * 100.0
	def _compute_plddt(
	plddt_per_atom: torch.Tensor | None,
	best_idx: int,
	n_tokens: int,
	atom_to_token: np.ndarray,
	atom_mask_bool: np.ndarray,
	) -> np.ndarray:
	"""Average the selected sample's per-atom pLDDT into per-token pLDDT.
	"""
	if plddt_per_atom is None:
	return np.full(n_tokens, 50.0, dtype=np.float32)
	per_atom = plddt_per_atom[best_idx].cpu().numpy()
	# Mean over the present atoms of each token; tokens with no atom stay at 0
	scored = atom_mask_bool & (atom_to_token < n_tokens)
	token_of_atom = atom_to_token[scored]
	totals = np.bincount(token_of_atom, weights=per_atom[scored], minlength=n_tokens)
	counts = np.bincount(token_of_atom, minlength=n_tokens)
	return (totals[:n_tokens] / np.maximum(counts[:n_tokens], 1)).astype(np.float32)

[ducta3141]

Suggested change

	def _compute_ptm(logits: torch.Tensor | None, n_tokens: int, has_frame: torch.Tensor | None = None) -> float:
	"""Compute predicted TM-score from PAE logits."""
	if logits is None:
	return float("nan")
	return _tm_score_from_pae_logits(logits, n_tokens, has_frame=has_frame)

[ducta3141]

Suggested change

	def _compute_iptm(
	logits: torch.Tensor | None,
	n_tokens: int,
	chain_indices: np.ndarray,
	has_frame: torch.Tensor | None = None,
	) -> float:
	"""Compute interface pTM from PAE logits (inter-chain pairs only)."""
	if logits is None:
	return float("nan")
	# Score only pairs that cross a chain boundary; a single chain has no interface.
	ci = torch.as_tensor(chain_indices, dtype=torch.long, device=logits.device)
	pair_mask = (ci.unsqueeze(-1) != ci.unsqueeze(-2)).to(dtype=torch.float32)
	if not bool(pair_mask.any()):
	return float("nan")
	return _tm_score_from_pae_logits(logits, n_tokens, pair_mask=pair_mask, has_frame=has_frame)

[ducta3141]

Suggested change

	probs = torch.softmax(logits.float(), dim=-1)
	n_bins = probs.shape[-1]
	# d0 = 1.24 * (max(N, 19) - 15)^(1/3) - 1.8, so the N floor of 19 keeps d0 > 0
	d0 = 1.24 * (max(n_tokens, 19) - 15) ** (1.0 / 3.0) - 1.8
	# Expected TM term per pair: E_bins[1 / (1 + (e_ij / d0)^2)]
	bin_centers = _bin_centers(0.0, 32.0, n_bins).to(device=probs.device)
	tm_per_bin = 1.0 / (1.0 + (bin_centers / d0) ** 2)
	tm_per_pair = (probs * tm_per_bin).sum(dim=-1) # (N, N)
	# Mean over the scored tokens j, for each aligned token i
	if mask is None:
	mask = torch.ones_like(tm_per_pair)
	n_scored = mask.sum(dim=-1) # (N,)
	tm_per_aligned = (tm_per_pair * mask).sum(dim=-1) / n_scored.clamp(min=1)
	# Maximum over the aligned tokens that are eligible and have something to score
	eligible = n_scored > 0
	if has_frame is not None:
	eligible = eligible & has_frame.to(device=eligible.device)
	if not bool(eligible.any()):
	return float("nan")
	return float(tm_per_aligned[eligible].max())

[ducta3141]

def _compute_pae(logits: torch.Tensor | None) -> np.ndarray | None:
    """Compute PAE matrix from PAE logits.

    Reducing the bin axis before the host copy makes the transfer ``n_bins``
    times smaller: an (N_token, N_token) matrix rather than the logits.
    """
    if logits is None:
        return None
    probs = torch.softmax(logits.float(), dim=-1)
    n_bins = probs.shape[-1]
    bin_centers = _bin_centers(0.0, 32.0, n_bins).to(device=probs.device)
    pae = (probs * bin_centers).sum(dim=-1).cpu().numpy()
    return np.round(pae, 3)

[ducta3141]

Also add some test cases for suggestions. And the MR should change to "match AF3 pTM/ipTM reduction and confidence bin centers and speedup postprocessor"