jpromera — JAX/Equinox weights for Promera

JAX/Equinox parameters for Promera, a dual-purpose biomolecular generative model for structure prediction and binder design. These weights were converted module-by-module from the official PyTorch checkpoint bjing-mit/promera (promera_2606.ckpt) and validated numerically against it (per-module relative error ~1e-6 vs PyTorch, true-fp32).

Use them with jpromera, the torch-free JAX/Equinox port of the model.

Files

file	what
promera_2606.eqx	Equinox-serialized parameters (eqx.tree_serialise_leaves)
promera_2606.skeleton.pkl	pytree skeleton (jax.ShapeDtypeStruct placeholders) so the model loads with no PyTorch / Promera installed

Usage

from huggingface_hub import hf_hub_download
import jpromera

# both files land in the same snapshot dir; load_model wants the shared stem
stem = hf_hub_download("escalante-bio/jpromera", "promera_2606.eqx")[:-4]
hf_hub_download("escalante-bio/jpromera", "promera_2606.skeleton.pkl")

jp = jpromera.serialize.load_model(stem)        # torch-free
out = jp.fold(feats, recycling_steps=4)         # trunk
coords, *_ = jp.sample(feats, out, num_steps=200,
                       diffusion_cfg=jpromera.DIFFUSION,
                       key=__import__("jax").random.PRNGKey(0))

Install jpromera (core is torch-free):

pip install git+https://github.com/escalante-bio/jpromera.git

License

MIT. Converted from bjing-mit/promera (© 2026 Bowen Jing and Mihir Bafna, MIT); the JAX/Equinox conversion is © 2026 Escalante Bio. The upstream copyright notice is preserved in the jpromera repository's LICENSE.