nboyd

name	pytorch-to-equinox
description	Translate a PyTorch nn.Module model to JAX/Equinox. Use when asked to convert, port, or translate PyTorch code to JAX, or when adding new modules to an existing PyTorch-to-Equinox codebase.

PyTorch to JAX/Equinox Translation Guide

You are translating PyTorch nn.Module code into JAX/Equinox modules. This is a manual, module-by-module translation — not an automatic transpiler. The goal is numerically equivalent forward passes with idiomatic JAX.

Read https://docs.kidger.site/equinox/all-of-equinox/ before starting. This is the authoritative reference for Equinox concepts (modules, pytrees, filtering, JIT, serialization).

nboyd

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "mosaic @ git+https://github.com/escalante-bio/mosaic.git",
#     "gemmi>=0.6.5",
#     "jax[cuda12]",
#     "numpy",
# ]
#
# [tool.uv]

nboyd

"""
Finetuning BoltzGen's structure module on hallucinated binder designs.
Simplified gist.
"""

import equinox as eqx
import jax
import jax.numpy as jnp
from joltzgen import AtomDiffusion, weighted_rigid_align
from mosaic.models.boltzgen import Sampler

nboyd

from math import sqrt

import equinox as eqx
import gemmi
import jax
import jax.numpy as jnp
import mosaic.losses.structure_prediction as sp
from mosaic.common import TOKENS
from mosaic.losses.protein_mpnn import InverseFoldingSequenceRecovery
from mosaic.losses.transformations import NoCys

nboyd

import modal

def download_boltz2():
    from mosaic.models.boltz2 import Boltz2
    Boltz2()


### Build modal image: install mosaic + deps and download boltz2 model.
image = (
    modal.Image.debian_slim(python_version="3.12")

nboyd

#########
#
# Low-hanging fruit:
#   1. Epitope selection! E.g. using Pesto-style predictor or something simple (SASA + charge?)
#   2. Cropping! Could be sped up by a factor of 5-10 by cropping the target properly
#   3. Early stopping. Kill runs that are going poorly in early stages
#   4. Filtering + ranking with more models
#   5. Different design models, hyperparameter tuning etc. E.g. Boltz2 + templates
#
import modal

nboyd

class FunctionWrapperContextManager:
    def __init__(self, module, function_names, callback):
        """
        Create a context manager that wraps multiple functions in a module.
        Calls to those functions are dispatched to a callback.

        Args:
            module (module): The module containing the functions to wrap
            function_names (list): List of function names to wrap