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.

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)

Requirements

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

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