Accelerate protein sequence design with JAX-optimized MPNN

Aminx: A functional interface to ProteinMPNN in JAX

Aminx is a JAX/Equinox reimplementation of the LigandMPNN codebase. It reproduces the PyTorch reference to ≥ 0.999 Pearson correlation across all five decoding paths, and runs 8-61× faster on a single structure (H200) by trading eager dispatch for jit/vmap/scan kernels.

What you get:

• A functional sample() / score() API - no model objects to wire up, no inference loop to write.
• A composable inference layer (StageSet) for swapping logit transforms, encode paths, and decode variants without touching kernel math.
• Numerical parity with upstream LigandMPNN, validated across unconditional, conditional, autoregressive, membrane, and side-chain-packer paths.
• Native batching across temperatures, backbones, and sequence lengths - operations you'd normally write as Python loops are compiled into single JAX kernels, so there's no recompilation penalty and no padding waste.

The batching point is worth unpacking: in vanilla JAX, running N temperatures requires either a Python loop (N separate compiled calls, or worse, N retraces) or manually writing a vmap. aminx has already done that work. Passing temperature=[0.1, 0.3, 0.7] vmaps over the temperature axis in one compiled call; scoring a mixed-length library reuses a compiled kernel per length bucket rather than recompiling per structure. The speedups in the tables below come from applying this pattern consistently to the operations most commonly written as loops in protein design workflows.

Performance

Benchmarked on H200 (NVIDIA SXM5), A100 (PCIe), L40s, and Blackwell (SM120). All figures are warm-call medians; see the full benchmark report for every hardware and configuration.

H200 - single-structure latency (seq_len=76)

Advanced Capability Benchmarks

These are capabilities ProteinMPNN doesn't expose, measured against the PyTorch baseline doing the equivalent work.

The 8× floor holds across hardware; ceilings reach 84-91× depending on the operation (A100: 8-91×, L40s: 8-85×, Blackwell SM120: 8-84×).

Documentation

Complete Documentation →

• Composition Guide - StageSet, InferencePlan, and the five extension points
• Parity Validation - numerical parity report vs the LigandMPNN reference

Validation

Aminx is validated against the upstream LigandMPNN reference (which includes ProteinMPNN behavior):

Full parity suite: 30/30 parity_heavy tests pass on the Engaging cluster (job 14203624). 575 fast tests pass locally (575 passed, 6 skipped, 2 xfailed).

Canonical parity docs (source of truth):

• Parity report (HTML)
• Parity report (PDF)

Root-level parity stubs are non-canonical; use the links above.

Quick Start

Installation

aminx uses uv as its package manager. Install uv first if you don't have it:

curl -LsSf https://astral.sh/uv/install.sh | sh

Then install aminx with your target hardware extra:

uv sync --extra cuda  # GPU (CUDA)
uv sync --extra tpu   # TPU
uv sync --extra cpu   # CPU-only

If you prefer pip:

pip install "aminx[cuda]"   # GPU
pip install "aminx[cpu]"    # CPU-only

Install as a standalone CLI tool (no virtual environment required):

uv tool install aminx
aminx spec validate run_spec.json

Or invoke without installing at all:

uvx aminx spec validate run_spec.json

High-level API

from aminx.io.weights import load_model
from aminx.run import sample, score, SamplingSpecification, ScoringSpecification

model = load_model(model_version="v_48_020", model_weights="original")

# --- Sampling ---
spec = SamplingSpecification(
    inputs="path/to/structure.pdb",
    num_samples=10,
    temperature=0.1,
    random_seed=42,
)
results = sample(spec)
sequences = results["sequences"]   # (num_samples, seq_len)
logits    = results["logits"]      # (num_samples, seq_len, 21)

# --- Scoring ---
spec = ScoringSpecification(
    inputs="path/to/structure.pdb",
    sequences_to_score=["MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEK"],
    temperature=1.0,
)
results = score(spec)
scores = results["scores"]         # negative log-likelihood per sequence

Composable Inference API

For control over fusion strategy, encode path, and decode variant without touching kernel math:

import jax.numpy as jnp
from aminx.host.plan import make_inference_plan, InferencePlan, InferenceComponents
from aminx.inference.encode import make_encode_fn
from aminx.inference import driver
from aminx.inference.logits import GeometricMeanLogits, ARLogitFuse
from aminx.run import SamplingSpecification
from aminx.types.stages import StageSet

# Option A: factory (resolves stages from spec automatically)
spec = SamplingSpecification(
    inputs="structure.pdb",
    num_samples=10,
    multi_state_strategy="geometric_mean",
    state_weights=[1.0, 0.8, 0.6],
)
plan = make_inference_plan(model, spec)

# Option B: manual assembly for full control
state_weights = jnp.array([1.0, 0.8, 0.6])
stage_set = StageSet(
    logit_transform=GeometricMeanLogits(weights=state_weights, temperature=1.2),
    ar_logit_transform=ARLogitFuse(),
)
plan = InferencePlan(
    model=model,
    components=InferenceComponents(
        encode_fn=make_encode_fn(model, use_rolling_state=False),
        driver=driver.decode,
        stage_set=stage_set,
    ),
)

result = plan.sample(bundle, key, config)   # → SampleResult(sequence, logits)
logits = plan.score(bundle, key, config)    # → (L, 21)

Plain callables and lambdas work in all StageSet slots (the driver uses eqx.filter_jit). Use eqx.Module only when the callable carries JAX array leaves (e.g. weights that need grad).

See the Composition Guide for the five extension points (logit_transform, ar_logit_transform, decode_step, sample_step, tie_group_fuse).

Advanced Capabilities

Temperature Array Sweep

Pass a list of temperatures to run all M in a single JIT-compiled call. The kernel vmaps over the temperature dimension, so per-temperature cost scales near-ideally.

from aminx.run import sample, SamplingSpecification

spec = SamplingSpecification(
    inputs="structure.pdb",
    temperature=[0.1, 0.3, 0.7, 1.0],  # M=4 temperatures, one JIT call
    num_samples=10,
)
results = sample(spec)
# One result dict per temperature, each with "sequences" and "logits"
for temp, result in zip(spec.temperature, results):
    print(f"T={temp}: {result['sequences'].shape}")

On H200 at M=8 (seq_len=76), per-temperature cost is 2.2 ms - the same total wall-clock as M=1 (17 ms) split across 8 temperatures. Measured scaling: 7.97-8.47× across H200/A100/L40s/Blackwell.

Deduplicated Scoring (DedupGather pattern)

When scoring a large ensemble where many sequences share a backbone, score only the K unique structures rather than all N. aminx JIT-caches a compiled kernel per length bucket, so K sequential plan.score() calls stay fast regardless of N.

from aminx.host.plan import make_inference_plan
from aminx.inference.bundle_builder import build_inference_bundle
from aminx.tiling.bucketing import BucketingConfig
import jax.random as random

plan = make_inference_plan(model, spec)
bucket_cfg = BucketingConfig()
key = random.PRNGKey(0)

# Build one bundle per unique backbone (K unique out of N total)
unique_bundles = []
for coords, mask, residue_index, chain_index, sequence in unique_structures:
    bundle, config = build_inference_bundle(
        coords=coords,
        mask=mask,
        residue_index=residue_index,
        chain_index=chain_index,
        sequence=sequence,
        ligand_coords=None,
        ligand_atom_types=None,
        ligand_mask=None,
        temperature=1.0,
        mode="score_conditional",
        inference=True,
        bucket_config=bucket_cfg,
    )
    unique_bundles.append((bundle, config))

# Score K unique structures; scatter the K scores to your N positions
scores = [plan.score(bundle, key, config) for bundle, config in unique_bundles]

At K=1/N=32 on H200, latency is 1.1 ms vs 92 ms for PyTorch scoring all 32 structures - an 80× speedup. The speedup is stable across K (80-82× from K=1 through K=32) because PyTorch's cost scales linearly with N while aminx's scales linearly with K.

Mixed-Length Batch Scoring

Score a library of proteins with different sequence lengths without padding waste. aminx rounds each length to the next power-of-2 boundary and reuses the JIT-compiled kernel across structures in the same bucket - one compile per bucket, not one per structure.

from aminx.host.plan import make_inference_plan
from aminx.inference.bundle_builder import build_inference_bundle
from aminx.tiling.bucketing import BucketingConfig
import equinox as eqx

plan = make_inference_plan(model, spec)
bucket_cfg = BucketingConfig()

# Build one bundle per structure — lengths can differ freely
bundles = []
for coords, mask, residue_index, chain_index, sequence in your_library:
    bundle, config = build_inference_bundle(
        coords=coords,
        mask=mask,
        residue_index=residue_index,
        chain_index=chain_index,
        sequence=sequence,
        ligand_coords=None,
        ligand_atom_types=None,
        ligand_mask=None,
        temperature=1.0,
        mode="score_conditional",
        inference=True,
        bucket_config=bucket_cfg,
    )
    bundles.append((bundle, config))

# Each call reuses the compiled XLA kernel for its length bucket
score_one = eqx.filter_jit(plan.score)
scores = [score_one(bundle, key, config) for bundle, config in bundles]

For a batch of [76, 150, 300, 500]-residue structures on H200, total latency is 4.1 ms versus 2280 ms for PyTorch (padded sequential) - a 554× speedup. The gain comes from two sources: JAX's compiled kernels (vs PyTorch eager dispatch) and avoiding padding to the longest sequence.

CLI

aminx ships a Typer CLI with four command groups: run, campaign, spec, and the spec emit-* family.

aminx run - run sampling and scoring pipelines

# Full sampling pipeline — constructs a RunSpecification and runs it end-to-end
aminx run sample \
  --inputs structure.pdb \
  --model-version v_48_020 \
  --model-weights original \
  --num-samples 10 \
  --random-seed 42

# Emit the spec JSON instead of running (useful for inspection or handoff)
aminx run sample --inputs structure.pdb --emit-json
aminx run sample --inputs structure.pdb --emit-json --out sample_spec.json

# score, jacobian, and inspect accept the same options;
# pass --emit-json to get the spec — the runner for these paths is not yet wired
aminx run score --inputs structure.pdb --sequences-to-score ACDEFGHIKLMNPQRSTVWY --emit-json

All four subcommands (sample, score, jacobian, inspect) share the same base option surface: --inputs, --model-weights, --model-version, --model-family, --batch-size, --backbone-noise, --random-seed, and the full RunSpecification field set. Spec construction failures exit 1; an unwired runner exits 2.

Heterogeneous inputs - local files and remote structures

The --inputs flag accepts multiple URI schemes, allowing you to mix local files and remote-fetched structures in a single command. Remote sources are fetched at CLI time to a local cache directory, making the resolved spec safely reproducible on cluster compute nodes (which are offline).

Accepted URI forms:

Example - mixed local and remote inputs:

# Fetch two RCSB structures and one AlphaFold structure; combine with a local directory
aminx run sample \
  --inputs pdb://1A3A \
  --inputs pdb://2ABC.pdb \
  --inputs afdb://P12345 \
  --inputs ./local_structures/ \
  --model-version v_48_020 \
  --num-samples 10

All remote sources are fetched to a cache directory (default: ~/.cache/aminx/inputs or $XDG_CACHE_HOME/aminx/inputs). Subsequent runs with the same accession reuse cached files and skip the network call entirely - the resolved spec is fully deterministic and cluster-safe.

Cache control and URI scheme override:

# Override cache directory
aminx run sample --inputs pdb://1A3A --cache-dir /tmp/aminx_cache --num-samples 10

# Force all --inputs entries to be treated as local paths (suppress scheme detection)
aminx run sample --inputs "pdb://1A3A" --input-type file --num-samples 10
# ^ Treats "pdb://1A3A" as a literal local filename, does not fetch

# Apply a default scheme to schemeless tokens
aminx run sample --inputs 1A3A --input-type pdb --num-samples 10
# ^ Treats "1A3A" as pdb://1A3A; fetches from RCSB

Error handling: Empty accessions (e.g., pdb://), unknown schemes (e.g., s3://, http://), and fetch failures raise clear errors naming the problematic entry. Use --inputs-fail-fast to exit on first error, or omit it to warn and skip individual entries.

aminx campaign - design campaign orchestration

# Create a campaign manifest from a base spec
aminx campaign plan \
  --inputs structures/target.pdb \
  --campaign-id pilot_v1 \
  --manifest-path pilot.manifest.json \
  --output-root outputs/pilot \
  --designs-per-library-type 50 \
  --samples-chunk-size 16

# Execute a single manifest row (used by distributed workers)
aminx campaign worker --manifest-path pilot.manifest.json --row-index 0

# Execute all rows; exits 1 if any row fails
aminx campaign run --manifest-path pilot.manifest.json

# Evaluate quality gates; exits 2 if the campaign is not promoted
aminx campaign gates --manifest-path pilot.manifest.json

# Plan a staged scale ramp
aminx campaign ramp-plan \
  --inputs structures/target.pdb \
  --campaign-id ramp_v1 \
  --manifest-dir ramp_manifests/ \
  --output-root outputs/ramp \
  --stage-designs-per-library-type 10,50,200 \
  --samples-chunk-size 16

# Evaluate ramp stage reports; exits 2 if not promoted
aminx campaign ramp-evaluate --report-path stage1.json --report-path stage2.json

--lock-backend distributed is not supported from the CLI (raises an error); use local_fs (the default).

aminx spec - validate and round-trip spec files

# Validate a spec JSON — exits 0 and prints "OK: <SpecClass>", or exits 1 with the parse error
aminx spec validate run_spec.json

# Round-trip a spec through the codec — prints the re-serialized JSON to stdout
aminx spec roundtrip run_spec.json
aminx spec roundtrip run_spec.json --out validated.json   # write to file instead
aminx spec roundtrip run_spec.json --compact              # single-line JSON

# Round-trip the portable subset (dict-serializable fields only, no JAX arrays)
aminx spec portable-roundtrip portable_spec.json
aminx spec portable-roundtrip portable_spec.json --compact

aminx spec emit-* - emit spec JSON without running

The emit-sample, emit-score, emit-jacobian, and emit-inspect subcommands are a convenience complement to aminx run <cmd> --emit-json - identical output, no runner invoked.

# Emit a sample spec JSON — equivalent to: aminx run sample ... --emit-json
aminx spec emit-sample \
  --inputs structure.pdb \
  --model-version v_48_020 \
  --model-weights original \
  --compact

aminx spec emit-sample --inputs structure.pdb --out sample_spec.json

Specs can also be serialized from Python with run_specification_to_json:

from aminx.run import run_specification_to_json
json_str = run_specification_to_json(spec)

Potts Model Family (aminx.potts)

The aminx.potts module provides a parallel model family integrating PottsMPNN with
Tree-Reweighted belief propagation (TRW) for global sequence design.

CLI: aminx potts run, aminx potts emit
Weight recapture: scripts/recapture/pottsmpnn_to_eqx.py
Architecture: Parallel model family - see .praxia/docs/decisions/260605_potts-parallel-not-stageset.md

Requirements

• Python ≥ 3.12
• JAX + Equinox (GPU/TPU/CPU via extras)
• uv sync --extra cpu for CPU-only; --extra cuda for GPU

Where model weights come from

The built wheel ships no checkpoints, so the aminx version pin does not by itself
determine which weights execute. Resolution order, and it stops at the first hit:

1. AMINX_WEIGHTS_DIR - when set, this directory is authoritative and fails closed: a
   checkpoint missing from it raises rather than quietly falling through to the Hub. Scope is
   checkpoint-id resolution only; an explicit local_path / --model-local-path still
   bypasses it by design, and logs a warning when it does.
2. Packaged resources (aminx/model_params/) - present in a source checkout, absent from
   the wheel.
3. Hugging Face Hub, pinned to HF_REVISION in aminx/io/weights.py.

Because the pin is a full commit SHA, a warm cache resolves with no network request at
all - measured at ~0.7 ms against ~236 ms for the unpinned form, which issued a HEAD on
every call. Pinning is faster and reproducible, not a tradeoff between them.

Recording which weights ran

from aminx.io.weights import weight_provenance

record = weight_provenance("proteinmpnn_v_48_020.eqx.zst")
print(record.source, record.sha256, record.hub_revision)

Store sha256 next to the aminx version in any result whose numbers depend on the weights -
the version alone does not identify them. weight_provenance shares one resolver with the
loader, so the record describes the file that actually executes for the checkpoint-id
route. It does not cover a run that passed local_path / --model-local-path, or one that
resolved through checkpoint_registry_path - those bypass resolution by design and log when
they do, and recording their weights is the caller's job.

Overriding the pin

Set AMINX_WEIGHTS_REVISION to reach a checkpoint newer than the pin, or to roll back:

export AMINX_WEIGHTS_REVISION=<commit-sha>

In a SLURM batch script, export it before the run line:

#SBATCH --job-name=aminx-score
export AMINX_WEIGHTS_REVISION=25fb7f6e985724dee7471c3bc18522fe33b9228e
uv run aminx run score --spec spec.json

The override is visible in WeightProvenance.hub_revision, which is read back from the
resolved cache path rather than from the variable - so an overridden run is still identifiable
from its own record. Setting either variable to a blank value is an error, not a request
for defaults: a blank value is almost always an unset variable interpolated into an
environment, and honouring it would change the weight source with no signal. Note this check
runs before the resolution order is known, so a blank AMINX_WEIGHTS_REVISION raises even on
a run that would have resolved from AMINX_WEIGHTS_DIR or packaged resources and never reached
the Hub. That is deliberate: a value that only fails once something happens to reach the Hub is
the failure mode this replaces.

Development

All five decoding paths are validated via parity_heavy tests - see Validation Reference below.

Architecture

aminx.run          ← SamplingSpecification, ScoringSpecification, sample(), score()
aminx.host.plan    ← InferencePlan, InferenceComponents, make_inference_plan()
aminx.types.stages ← StageSet (the composition interface)
aminx.inference    ← driver.decode, logits (LOGIT_STRATEGIES, TIE_GROUP_STRATEGIES)
aminx.model        ← LigandMPNN, Packer (Equinox modules, JIT-safe)
aminx.sampling     ← sample() kernel
aminx.scoring      ← score() kernel
aminx.cli          ← aminx run / campaign / spec (Typer entry point)

StageSet is the seam between the host layer and the JAX-traced kernels: everything above it is Python-land, everything below it is traced. See the Composition Guide.

Multiprocessing

Importing aminx does not set the multiprocessing start method. If your notebook or script spawns worker processes, call configure_multiprocessing() once at startup (see aminx.runtime); the campaign CLI does this for you.

Validation Reference

# Install project dependencies (CPU/dev/tests path)
uv sync --extra cpu --extra dev --extra tests --group dev
source .venv/bin/activate

# Checkout reference implementation (pinned commit used in CI)
git clone https://github.com/dauparas/LigandMPNN.git reference_ligandmpnn_clone
cd reference_ligandmpnn_clone && git checkout 3870631 && cd ..

# Optional strict preflight per parity tier
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/check_parity_prereqs.py --reference-path "$REFERENCE_PATH" --project-root . --tier parity_heavy
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/check_parity_prereqs.py --reference-path "$REFERENCE_PATH" --project-root . --tier parity_audit

# Validate parity asset cache/checksums
uv run python scripts/check_parity_assets.py --tier parity_fast
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/check_parity_assets.py --tier parity_heavy
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/check_parity_assets.py --tier parity_audit

# Run fast deterministic parity checks
uv run pytest tests/parity -m parity_fast -v

# Run reference-backed heavy parity checks
REFERENCE_PATH=./reference_ligandmpnn_clone \
  PRXTEIN_PARITY_TIER=parity_heavy \
  uv run pytest tests/parity tests/model/test_ligandmpnn_equivalence.py -m parity_heavy -v

# Convert full checkpoint families and run parity_audit checks
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/convert_parity_family_weights.py \
    --project-root . \
    --reference-path "$REFERENCE_PATH" \
    --tier parity_audit \
    --skip-existing
REFERENCE_PATH=./reference_ligandmpnn_clone \
  PRXTEIN_PARITY_TIER=parity_audit \
  uv run pytest tests/parity tests/model/test_ligandmpnn_equivalence.py -m parity_audit -v

# Collect expanded parity evidence (multi-backbone + synthetic random cases)
REFERENCE_PATH=./reference_ligandmpnn_clone \
  uv run python scripts/collect_parity_evidence.py \
    --project-root . \
    --case-corpus tests/parity/parity_case_corpus.json \
    --output-dir docs/parity/reports/evidence

# Render Markdown/HTML report and export PDF with embedded plots/tables
uv run python scripts/generate_parity_report.py --project-root . --output-dir docs/parity --pdf

AMINX_VERIFY=1 uv run pytest path/to/test.py -v