Power-schedule EDM initialization and sampler step

Your task

Function signature

edm_sampler_step(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    steps: int = 4,
    schedule: str = 'af3',
    cap_schedule: bool = False,
    align: bool = False,
) -> dict[str, torch.Tensor]

Worked example

One step without added noise

import torch

activations = {
    'x_t': torch.tensor([[[2.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'denoised': torch.tensor([[[0.0, 2.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'epsilon': torch.tensor([[[0.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'initial_noise': torch.tensor([[[1.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'sigma_min': torch.tensor(1.0, dtype=torch.float64, requires_grad=True),
    'sigma_max': torch.tensor(2.0, dtype=torch.float64, requires_grad=True),
    'rho': torch.tensor(1.0, dtype=torch.float64, requires_grad=True),
    'gamma_0': torch.tensor(0.0, dtype=torch.float64, requires_grad=True),
    'gamma_min': torch.tensor(0.0, dtype=torch.float64, requires_grad=True),
}
parameters = {
    'noise_scale': torch.tensor(1.0, dtype=torch.float64, requires_grad=True),
    'step_scale': torch.tensor(1.0, dtype=torch.float64, requires_grad=True),
}
masks = {
    'valid_mask': torch.tensor([[True]], dtype=torch.bool),
    'diffuse_mask': torch.tensor([[True]], dtype=torch.bool),
}
structure = {
    'step': torch.tensor([0], dtype=torch.int64),
}
steps = 1
schedule = 'af3'
cap_schedule = False
align = False

result = edm_sampler_step(activations=activations, parameters=parameters, masks=masks, structure=structure, steps=steps, schedule=schedule, cap_schedule=cap_schedule, align=align)

Expected return

{
    'schedule': torch.tensor([2.0, 1.0], dtype=torch.float64),
    'gamma': torch.tensor([0.0], dtype=torch.float64),
    'sigma_hat': torch.tensor([2.0], dtype=torch.float64),
    'x_initial': torch.tensor([[[2.0, 0.0, 0.0]]], dtype=torch.float64),
    'x_hat': torch.tensor([[[2.0, 0.0, 0.0]]], dtype=torch.float64),
    'euler_base': torch.tensor([[[2.0, 0.0, 0.0]]], dtype=torch.float64),
    'direction': torch.tensor([[[1.0, -1.0, 0.0]]], dtype=torch.float64),
    'x_next': torch.tensor([[[1.0, 1.0, 0.0]]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Construct a power schedule and one EDM update, with all random variates and denoiser outputs supplied. All real inputs have one common \mathrm{float32}/64 dtype.
• Scalar \sigma_{\mathrm{min}}>0, \sigma_{\mathrm{max}}\ge \sigma_{\mathrm{min}} and \rho>0 define s(r)=((1-r)\cdot \sigma_{\mathrm{max}}^{1/\rho}+r\cdot \sigma_{\mathrm{min}}^{1/\rho})^{\rho}. steps is a positive integer, default 4.
• schedule='af3' (default) returns s(i/\mathrm{steps}), i=0..steps: its final value is POSITIVE sigma_min, as in the inspected AF3 JAX release. 'esmfold2' returns steps positive levels s(i/(\mathrm{steps}-1)) followed by zero; \mathrm{steps}=1 instead returns [\sigma_{\mathrm{max}},0].
• cap_schedule=False by default; True is allowed only for esmfold2 and replaces the schedule by [\sigma_{\mathrm{cap}}] concatenated with every original \mathrm{level}\le \sigma_{\mathrm{cap}}, preserving order. \sigma_{\mathrm{cap}}>0 is then required; equality retains the original level too.
• For each integer step [b] in [0,\mathrm{len}(\mathrm{schedule})-2], let \mathrm{prev}=\mathrm{schedule}[\mathrm{step}], \mathrm{next}=\mathrm{schedule}[\mathrm{step}+1]. \gamma=\gamma_{0} if \mathrm{next}>\gamma_{\mathrm{min}}, otherwise zero (strict inequality); \sigma_{\mathrm{hat}}=\mathrm{prev}\cdot (1+\gamma).
• \gamma_{0}\ge 0, \gamma_{\mathrm{min}}\ge 0, \mathrm{noise\_scale}\ge 0 and \mathrm{step\_scale}\ge 0 are supplied scalars. x_{\mathrm{initial}}=\mathrm{schedule}[0]\cdot \mathrm{initial\_noise}; x_{\mathrm{hat}}=x_{t}+\mathrm{noise\_scale}\cdot \mathrm{prev}\cdot \sqrt{\gamma\cdot (2+\gamma)}\cdot \epsilon.
• Use a safe zero-churn branch before sqrt; inactive epsilon may be nonfinite and is ignored when \gamma=0. align=False normally sets \mathrm{euler\_base}=x_{\mathrm{hat}}.
• If True, alignment_R [B,3,3] and alignment_t [B,3] are supplied, and \mathrm{euler\_base}=R\cdot x_{\mathrm{hat}}+t in column-vector notation (row arrays multiply R^{\top}).
• These transforms are computed outside this exercise, for example by weighted_kabsch_alignment; the denoiser output denoised is already in the destination frame.
• \mathrm{direction}=(\mathrm{euler\_base}-\mathrm{denoised})/\sigma_{\mathrm{hat}}; x_{\mathrm{next}}=\mathrm{euler\_base}+\mathrm{step\_scale}\cdot (\mathrm{next}-\sigma_{\mathrm{hat}})\cdot \mathrm{direction}. An equivalent affine blend avoids unnecessarily large residuals.
• A zero next level reaches denoised only if \mathrm{step\_scale}=1; the common over-relaxed scale 1.5 need not. No second denoiser call, hidden callback or Heun corrector. Coordinates, denoised, epsilon and initial_noise are [B,N,3]; masks are [B,N]. valid_mask=True means observed; diffuse_mask=True means movable.
• Invalid atoms return zero in all coordinate outputs. Fixed valid atoms retain x_t in x_initial/x_hat/euler_base/x_next and have zero direction; their denoised/noise payloads are ignored. Sanitize masked payloads before arithmetic. B\ge 1,N\ge 0; all-masked examples are allowed.
• Active coordinates/variates and scalar parameters must be finite and yield representable results. No centering here: rigid_coordinate_augmentation is a separate operation. Invalid schedule/step bounds and nonpositive schedule scales raise ValueError.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def edm_sampler_step(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    steps: int = 4,
    schedule: str = 'af3',
    cap_schedule: bool = False,
    align: bool = False,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    schedule: (S,); uncapped S=steps+1, capped S=1+number of {s in
        uncapped:s<=sigma_cap}.
    gamma: (B,).
    sigma_hat: (B,).
    x_initial: (B, N, 3).
    x_hat: (B, N, 3).
    euler_base: (B, N, 3).
    direction: (B, N, 3).
    x_next: (B, N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement edm_sampler_step")

Sources and further reading

• www.nature.com
• google-deepmind/alphafold3 · diffusion_head.py
• biohub.ai
• huggingface/transformers · generation_esmfold2.py