[mahdi-shafiei]

Summary

 _compute_differentiable_fp (CircularFingerprintOperator, MACCSKeysOperator)
   node_hidden = message_passing(...)
-  node_hidden = node_hidden * node_mask[:, None]     # zero rows, then...
   logits = biased layers(node_hidden)                # ...bias makes them nonzero
   bits = soft_ops.greater(logits, 0)
-  fingerprint = soft_ops.max(bits, axis=0)           # padding atoms in the max
+  fingerprint = soft_ops.max(bits, axis=0, where=node_mask > 0)

Two commits:

1. soft_ops.max and soft_ops.argmax gain where, with the semantics of jax.nn.standardize(where=) and jax.nn.softmax(where=): an excluded entry takes no part in the standardization statistics and gets zero weight and zero gradient. Smooth-mode proj_simplex is jax.nn.softmax(x / softness), so the masked form is JAX's own. It's supported for mode="hard" and for method="softsort" in smooth mode. neuralsort, the sorting network, optimal transport and the sparse simplex modes refuse where, because no implementation found (softjax, fast-soft-sort, OTT, the SoftSort and NeuralSort code) has a masked form of them.
2. The two fingerprint readouts use it.

Evidence

Before, padded-and-masked against unpadded output for every drug-discovery graph operator (ethanol and benzene padded to 13 atoms):

operator	before	after
CircularFingerprint	0.135	1.8e-7
MACCSKeys	0.058	8.9e-8
AttentiveFP (already masks its scores)	0.0	unchanged
Property / ADMET predictors (sum and mean readouts)	~1e-6	unchanged

The remaining 1e-7 is float32 summation order over a longer axis, not padding.

Tests, written first:

• tests/core/test_soft_ops/test_masked_max.py (8). Inputs whose means, variances and softmax denominators are exact under any summation order, so the comparisons are exact: masked soft argmax, soft max, per-row masks, hard mode, zero gradient for excluded entries, and refusal for unsupported methods. Before: 8 failed (no where). After: 8 passed.
• tests/operators/drug_discovery/test_padded_readout.py (2). Padded versus unpadded methanol, exact by construction: parameters from {-1/4, 0, 1/4}, one-hot features, and two heavy atoms, so every sum over atoms has at most two nonzero terms. Against main's code: 15 of 16 bits differ in both operators. After: equal bit for bit.
• An earlier draft of this test, "the fingerprint ignores what the padding atoms hold", passed on main too. Main already zeroes the padding rows; the leak is the biased layers' constant output. It was replaced by the comparison above.

Gate: the soft-op, drug-discovery, DTI and MolNet benchmark tests pass 587 of 587 on CPU. The first commit was gated on its own exported index (511 soft-op and drug-discovery tests). pre-commit run --all-files passes all 24 hooks.

Merge Danger

Door: two-way.

Blast Radius: padded fingerprints.

• What changes: fingerprints computed on padded input, which means every DTI drug encoding (dti.py indexes padded batch rows) and any user code that pads. Unpadded fingerprints don't change.
• New keyword: where is keyword-only, and every existing soft_ops.max and argmax call works as before.
• Not in this PR: protein/secondary_structure.py:367-371 zeroes local H-bonds before a soft max, the same defect class, measured at up to 0.069 on a random map. It's a different operator, and a residue with no non-local partner needs handling, so it gets its own PR.