Upload 83 files

README.md

@@ -0,0 +1,213 @@
+# pCoMole
+
+**pCoMole** (Pareto-Constrained Molecule editing) steers a pretrained Edit Flow toward user-specified objectives while enforcing hard terminal constraints. This repository is a minimal, clone-and-run package for:
+
+1. Training, evaluating, and sampling **Edit Flows**
+2. Running **pCoMole** multi-objective editing on **GFP** and **peptidomimetics**
+
+Paper: *pCoMol: Pareto-Constrained Molecule Editing with Discrete Flows*
+
+The original working tree is left untouched. This folder is the public-facing copy.
+
+## Repository layout
+
+```text
+pCoMole/
+  train.py / evaluate.py          # Edit Flow train / val-loss eval
+  configs/                        # GFP (protein) and SELFIES configs
+  model/ logic/ flow_matching/    # Edit Flow implementation
+  smiles_tokenizer/               # SMILES SPE + SELFIES vocab
+  data/selfies/28k_mimetics/      # shipped SELFIES training split
+  gfp/                            # GFP generation + pCoMole
+    pcomole.py                    # official GFP editor
+    generate.py
+    ckpt/last_2.ckpt
+    classifier_ckpt/best.pt
+    FPredX/                       # excitation / brightness / emission models
+  peptidomimetics/                # peptidomimetic generation + pCoMole
+    pcomole.py
+    generate.py
+    objectives.py                 # PeptiVerse + Admetica/DeepDTAGen mix
+    ckpt/SELFIES_EditFlows.ckpt
+    ckpt/SMILES_BindEvaluator.ckpt
+    ckpt/admetica/                # LD50, solubility, Caco-2, half-life
+    ckpt/deepdtagen/
+```
+
+## 1. Environment
+
+A CUDA GPU is strongly recommended.
+
+```bash
+git clone <this-repo-url> pCoMole
+cd pCoMole
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+HuggingFace model downloads happen on first use (`facebook/esm2_t33_650M_UR50D`, `aaronfeller/PeptideCLM-23M-all`, ChemBERTa).
+
+### PeptiVerse (required for peptidomimetic pCoMole)
+
+Peptidomimetic property oracles use the latest [PeptiVerse](https://huggingface.co/ChatterjeeLab/PeptiVerse) SMILES predictors. Clone it **next to** this repo, or set `PEPTIVERSE_ROOT`:
+
+```bash
+# recommended: sibling of pCoMole
+cd ..
+git clone https://huggingface.co/ChatterjeeLab/PeptiVerse
+cd PeptiVerse
+pip install -r requirements.txt
+# follow PeptiVerse README to download model weights
+```
+
+pCoMole looks for PeptiVerse in this order:
+
+1. `$PEPTIVERSE_ROOT`
+2. `../PeptiVerse` relative to this repository
+
+PeptiVerse scores the peptide-like side of each candidate. Admetica + DeepDTAGen still score the small-molecule side, and the original peptide-likeness mix logic is unchanged.
+
+### MAFFT (required for GFP pCoMole)
+
+GFP excitation, brightness, and emission oracles run through FPredX and need [MAFFT](https://mafft.cbrc.jp/alignment/software/) on your `PATH`.
+
+```bash
+# conda
+conda install -c bioconda mafft
+
+# or point to a local binary
+export MAFFT_PATH=/path/to/mafft
+```
+
+Confirm with `mafft --version` before running GFP editing.
+
+## 2. Train, evaluate, and sample Edit Flows
+
+Run every command from the **repository root**.
+
+### SELFIES peptidomimetic Edit Flow
+
+Training data is shipped at `data/selfies/28k_mimetics`.
+
+```bash
+# train
+python train.py --config configs/config_selfies.yaml
+# optional: python train.py --config configs/config_selfies.yaml --wandb
+
+# evaluate a checkpoint on the validation split
+python evaluate.py \
+  --config configs/config_selfies.yaml \
+  --ckpt peptidomimetics/ckpt/SELFIES_EditFlows.ckpt
+
+# unconditional / seeded generation
+python peptidomimetics/generate.py \
+  --config configs/config_selfies.yaml \
+  --ckpt peptidomimetics/ckpt/SELFIES_EditFlows.ckpt \
+  --input 'CSCC[C@H](NC(=O)[C@@H]1CCCN1C(=O)[C@H](Cc1ccc(O)cc1)NC(=O)[C@H](CCCNC(=N)N)NC(=O)[C@H](CO)NC(=O)[C@H](Cc1ccc(O)cc1)NC(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H](N)CO)C(=O)N[C@@H](CC(=O)O)C(=O)O' \
+  --num_steps 30 \
+  --num_samples 2 \
+  --output_csv outputs/peptidomimetic_unconditional.csv
+```
+
+Or: `bash peptidomimetics/scripts/train.sh` and `bash peptidomimetics/scripts/generate.sh`.
+
+### GFP protein Edit Flow
+
+A trained GFP checkpoint is shipped at `gfp/ckpt/last_2.ckpt`. The original GFP training arrows are **not** included. To retrain, put a HuggingFace `load_from_disk` dataset under `data/gfp/{train,validation}` (or edit `configs/config_gfp.yaml`).
+
+```bash
+python evaluate.py --config configs/config_gfp.yaml --ckpt gfp/ckpt/last_2.ckpt
+
+python gfp/generate.py \
+  --config configs/config_gfp.yaml \
+  --ckpt gfp/ckpt/last_2.ckpt \
+  --input 'MSSGALLFHGKIPYVVEMEGNVDGHTFSIRGKGYGDASVGKVDAQFICTTGDVPVPWSTLVTTLTYGAQCFAKYGPELKDFYKSCMPDGYVQERTITFEGDGNFKTRAEVTFENGSVYNRVKLNGQGFKKDGHVLGKNLEFNFTPHCLYIWGDQANHGLKSAFKICHEITGSKGDFIVADHTQMNTPIGGGPVHVPEYHHMSYHVKLSKDVTDHRDNMSLKETVRAVDCRKTYDFDAGSGDTS' \
+  --num_steps 10 \
+  --num_samples 8 \
+  --output_csv outputs/gfp_unconditional.csv
+```
+
+## 3. pCoMole multi-objective editing
+
+### GFP
+
+Official editor: `gfp/pcomole.py` (length + excitation + brightness, with GFP-classifier and emission constraints).
+
+```bash
+bash gfp/scripts/pcomole.sh
+```
+
+Equivalent command:
+
+```bash
+python gfp/pcomole.py \
+  --root_dir gfp/FPredX \
+  --config configs/config_gfp.yaml \
+  --ckpt gfp/ckpt/last_2.ckpt \
+  --num_steps 10 \
+  --num_candidates 50 \
+  --num_rollouts 10 \
+  --objective_weights 3 1 1 \
+  --output_file outputs/gfp_length_excitation_brightness.csv \
+  --input 'MSSGALLFHGKIPYVVEMEGNVDGHTFSIRGKGYGDASVGKVDAQFICTTGDVPVPWSTLVTTLTYGAQCFAKYGPELKDFYKSCMPDGYVQERTITFEGDGNFKTRAEVTFENGSVYNRVKLNGQGFKKDGHVLGKNLEFNFTPHCLYIWGDQANHGLKSAFKICHEITGSKGDFIVADHTQMNTPIGGGPVHVPEYHHMSYHVKLSKDVTDHRDNMSLKETVRAVDCRKTYDFDAGSGDTS'
+```
+
+If the output filename contains `length_excitation_brightness`, `length_brightness`, `length_excitation`, or `length`, the matching objective subset is used. Any other name uses all three objectives.
+
+### Peptidomimetics
+
+Official editor: `peptidomimetics/pcomole.py`.
+
+Objectives (7 scores when `--specificity` is on):
+
+1. non-toxicity
+2. solubility
+3. permeability
+4. half-life
+5. affinity
+6. motif
+7. specificity
+
+PeptiVerse provides the peptide-side scores. Admetica + DeepDTAGen provide the small-molecule side. The two are mixed by peptide-likeness, same as the paper code. Motif / specificity still use the shipped BindEvaluator checkpoint.
+
+```bash
+# after PeptiVerse is cloned and its weights are downloaded
+bash peptidomimetics/scripts/pcomole.sh
+```
+
+`--target` and `--motifs` are required for affinity and motif oracles. `--specificity` adds the seventh objective.
+
+## 4. Included checkpoints
+
+| Path | Role |
+|---|---|
+| `gfp/ckpt/last_2.ckpt` | GFP Edit Flow |
+| `gfp/classifier_ckpt/best.pt` | GFP hard constraint |
+| `gfp/FPredX/{ex,bright,em}_model` | FPredX property models |
+| `peptidomimetics/ckpt/SELFIES_EditFlows.ckpt` | peptidomimetic Edit Flow |
+| `peptidomimetics/ckpt/SMILES_BindEvaluator.ckpt` | motif / specificity |
+| `peptidomimetics/ckpt/admetica/*.ckpt` | Admetica blending oracles |
+| `peptidomimetics/ckpt/deepdtagen/` | DeepDTAGen affinity |
+
+Together these are about 10 GB. HuggingFace uploads should use Git LFS for the `.ckpt` / `.pt` / `.pth` files.
+
+## 5. Environment variables
+
+| Variable | Purpose |
+|---|---|
+| `PEPTIVERSE_ROOT` | PeptiVerse repo root (manifest + `training_classifiers/`) |
+| `MAFFT_PATH` | MAFFT binary, if it is not on `PATH` |
+| `GFP_CLASSIFIER_CKPT` | optional override for the GFP classifier |
+
+## 6. Notes
+
+- Run scripts from the repo root, or use the wrappers under `*/scripts/`.
+- GFP FPredX will fail immediately if `mafft` is missing; that is expected.
+- Peptidomimetic pCoMole will fail at oracle init if PeptiVerse is missing or its weights were not downloaded.
+- Cas9 is not part of this release.
+- Training writes to `outputs/` by default.
+
+## Citation
+
+If you use this code, please cite the pCoMole paper and [PeptiVerse](https://www.biorxiv.org/content/10.64898/2025.12.31.697180).

configs/config_gfp.yaml

@@ -0,0 +1,75 @@
+task: protein
+
+compute:
+  ngpus: 1
+  nodes: 1
+
+logging:
+  enable_wandb: False
+  project: pCoMole
+
+data:
+  # GFP training arrows are not shipped. Point these at your own
+  # HuggingFace load_from_disk() dataset if you want to retrain.
+  train_path: data/gfp/train
+  val_path: data/gfp/validation
+  num_workers: 1
+
+training:
+  batch_size: None
+  snapshot: 100
+  eval_freq: 1000
+  perplexity_freq: 100
+  seed: 42
+  grad_accumulation_steps: 2
+  loc_prop_path: False
+  lam_prop: 1
+  reparameterize: True
+  gamma_rate: 1
+  gamma_edit: 1
+  use_aux_ce: False
+  aux_ce_weight: 1.0
+  no_training: False
+
+eval:
+  batch_size: None
+  sample_batch_size: 4
+  perplexity: False
+  perplexity_batch_size: 4
+
+optim:
+  weight_decay: 0.05
+  optimizer: AdamW
+  lr: 1.0e-6
+  beta1: 0.9
+  beta2: 0.95
+  eps: 1.0e-8
+  grad_clip: 1.0
+  eta_min_ratio: 0.1
+  fused: false
+  n_epochs: 1000
+  warmup_ratio: 0.1
+
+flow:
+  source_distribution: uniform
+  loss_function: editflows
+  exponent: 2.0
+  scheduler_type: polynomial
+  sampling_steps: 256
+  eps_id: -1
+
+model:
+  d_model: 768
+  max_len: 2048
+  n_layers: 12
+  n_heads: 12
+  dropout: 0.1
+  attn_dropout: 0.0
+  proj_dropout: 0.0
+  compile: False
+  esm_model_name: facebook/esm2_t33_650M_UR50D
+  freeze_esm: true
+  p_optimal: 0
+  scale_size: 0.5
+
+work_dir: outputs/gfp

configs/config_selfies.yaml

@@ -0,0 +1,57 @@
+task: selfies
+
+compute:
+  ngpus: 1
+  nodes: 1
+
+data:
+  train_path: data/selfies/28k_mimetics/train
+  val_path: data/selfies/28k_mimetics/validation
+  num_workers: 4
+
+training:
+  batch_size: 4
+  snapshot: 100
+  eval_freq: 1000
+  perplexity_freq: 100
+  seed: 42
+
+eval:
+  batch_size: 4
+  sample_batch_size: 4
+  perplexity: False
+  perplexity_batch_size: 4
+
+optim:
+  weight_decay: 0.03
+  optimizer: AdamW
+  lr: 3.0e-4
+  beta1: 0.9
+  beta2: 0.95
+  eps: 1.0e-8
+  grad_clip: 1.0
+  eta_min_ratio: 0.1
+  fused: false
+  n_epochs: 100
+  warmup_ratio: 0.1
+
+flow:
+  source_distribution: uniform
+  loss_function: editflows_reparam
+  exponent: 2.0
+  scheduler_type: polynomial
+  sampling_steps: 256
+  eps_id: -1
+
+model:
+  hidden_size: 768
+  cond_dim: 128
+  length: 2048
+  n_blocks: 8
+  n_heads: 12
+  dropout: 0.1
+  compile: False
+  p_optimal: 0
+  scale_size: 1.0
+
+work_dir: outputs/selfies

data/selfies/28k_mimetics/dataset_dict.json

@@ -0,0 +1 @@
+{"splits": ["train", "validation", "test"]}

data/selfies/28k_mimetics/test/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d709f3386113ebc6abe9400950274e2e41aaf7063b31e67d6ba15f3239462251
+size 786928

data/selfies/28k_mimetics/test/dataset_info.json

@@ -0,0 +1,36 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int8",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "seq_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "batch_size": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

data/selfies/28k_mimetics/test/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "98db6912c1e59254",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

data/selfies/28k_mimetics/tokenizer/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "type": "selfies-wordlevel",
+  "version": "2.1.1",
+  "special_tokens": {
+    "pad": "<pad>",
+    "bos": "<bos>",
+    "eos": "<eos>"
+  },
+  "semantic_constraints": null
+}

data/selfies/28k_mimetics/tokenizer/vocab.json

@@ -0,0 +1,46 @@
+{
+  "<bos>": 1,
+  "<eos>": 2,
+  "<pad>": 0,
+  "[#Branch1]": 3,
+  "[#Branch2]": 4,
+  "[#C]": 5,
+  "[#N]": 6,
+  "[/-Ring1]": 7,
+  "[/C]": 8,
+  "[/N]": 9,
+  "[=Branch1]": 10,
+  "[=Branch2]": 11,
+  "[=C]": 12,
+  "[=N]": 13,
+  "[=O]": 14,
+  "[=Ring1]": 15,
+  "[=Ring2]": 16,
+  "[=S]": 17,
+  "[Br]": 18,
+  "[Branch1]": 19,
+  "[Branch2]": 20,
+  "[C@@H1]": 21,
+  "[C@@]": 22,
+  "[C@H1]": 23,
+  "[C@]": 24,
+  "[CH1-1]": 25,
+  "[C]": 26,
+  "[Cl]": 27,
+  "[F]": 28,
+  "[I]": 29,
+  "[N+1]": 30,
+  "[N-1]": 31,
+  "[NH1]": 32,
+  "[N]": 33,
+  "[O-1]": 34,
+  "[O]": 35,
+  "[P]": 36,
+  "[Ring1]": 37,
+  "[Ring2]": 38,
+  "[S]": 39,
+  "[\\-Ring1]": 40,
+  "[\\-Ring2]": 41,
+  "[\\C]": 42,
+  "[\\N]": 43
+}

data/selfies/28k_mimetics/train/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d7e3d408b36d7c8a75b0d36f20d09b744c1c845961401c72a4023a9a6ffa3f6
+size 6084456

data/selfies/28k_mimetics/train/dataset_info.json

@@ -0,0 +1,36 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int8",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "seq_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "batch_size": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

data/selfies/28k_mimetics/train/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "58b5e20848e4941d",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

data/selfies/28k_mimetics/validation/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6667b7a12a440b50d996c8523813df0965a718e1a467ef2162fdea65d9b972cb
+size 778640

data/selfies/28k_mimetics/validation/dataset_info.json

@@ -0,0 +1,36 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int8",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "seq_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "batch_size": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

data/selfies/28k_mimetics/validation/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "8dc499c569440f6b",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

evaluate.py

@@ -0,0 +1,39 @@
+import argparse
+
+import torch
+import pytorch_lightning as pl
+
+from paths import add_repo_to_sys_path, resolve_path
+from train import build_dataloaders, build_editflow, load_config
+
+add_repo_to_sys_path()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate an Edit Flow checkpoint on the validation split")
+    parser.add_argument("--config", type=str, required=True)
+    parser.add_argument("--ckpt", type=str, required=True)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    editflow, _, _, _, _, _, _ = build_editflow(cfg, device=device)
+
+    ckpt = torch.load(resolve_path(args.ckpt), map_location=device, weights_only=False)
+    state = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
+    editflow.load_state_dict(state, strict=False)
+    editflow.eval()
+
+    _, val_dataloader = build_dataloaders(cfg)
+    trainer = pl.Trainer(
+        accelerator="gpu" if torch.cuda.is_available() else "cpu",
+        devices=1,
+        logger=False,
+        enable_checkpointing=False,
+    )
+    metrics = trainer.validate(editflow, val_dataloader, verbose=True)
+    print(metrics)
+
+
+if __name__ == "__main__":
+    main()

flow_matching/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+__version__ = "1.0.10"

flow_matching/loss/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .generalized_loss import MixturePathGeneralizedKL
+from .editflows_loss import EditFlowsLoss, EditFlowsLossReParam
+
+__all__ = [
+    "MixturePathGeneralizedKL",
+    "EditFlowsLoss",
+    "EditFlowsLossReParam"
+]

flow_matching/loss/editflows_loss.py

@@ -0,0 +1,320 @@
+# Copyright (c) Meta Platforms, Inc.
+# All rights reserved.
+
+from typing import List, Optional, Sequence
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.modules.loss import _Loss
+
+# If you kept the adapter around only for scheduler access you can drop it entirely,
+# because we pass precomputed_weight from training.step. Keeping it is harmless though.
+class EditFlowsLoss(_Loss):
+    """
+    Edit Flows loss (Eq. 23), ragged version:
+      L_i =  (sum_j λ_ins[i][j] + sum_j λ_del[i][j] + sum_j λ_sub[i][j])
+            - w_i * ( sum_{ins events e} [log λ_ins[i][slot_e] + log Q_ins[i](y_e)]
+                    + sum_{del j}       [log λ_del[i][j]]
+                    + sum_{sub j}       [log λ_sub[i][j] + log Q_sub[i](y_j)] )
+    """
+
+    def __init__(self, reduction: str = "mean") -> None:
+        super().__init__(None, None, reduction)
+
+    def forward(
+        self,
+        lam_ins: torch.Tensor,     # (B, L)
+        logits_ins: torch.Tensor,  # (B, L, V)
+        lam_del: torch.Tensor,     # (B, L)
+        lam_sub: torch.Tensor,     # (B, L)
+        logits_sub: torch.Tensor,  # (B, L, V)
+        z_t: torch.Tensor,         # (B, N) aligned, with eps_id
+        z_1: torch.Tensor,         # (B, N) aligned target, with eps_id
+        x_t: torch.Tensor,         # (B, L)
+        valid_mask: torch.Tensor,  # (B, L) bool, False==padding, True==valid
+        precomputed_weight: torch.Tensor,  # (B,) or ()  = kappa_dot/(1-kappa)
+        eps_id: int,
+        bos_id: int,
+        eos_id: int,
+    ) -> torch.Tensor:
+        """
+        Implements Eq. 23 style loss for Edit Flows.
+
+        We:
+        1) penalize total outgoing rate on x_t
+        2) for every column in (z_t, z_1) that still differs, we map it to an x_t edit
+            and add  - w * log u_required
+
+        BOS/EOS handling:
+        - z_t and z_1 are aligned and already contain BOS/EOS
+        - we do NOT allow edits that delete/replace BOS/EOS
+        - we also skip columns whose target token is BOS/EOS (they should already match)
+        """
+        device = lam_ins.device
+        B, L = x_t.shape
+        _, N = z_t.shape
+
+        # 1. ----- RATE TERM -----
+        # valid_mask: True = real token, False = pad
+        valid_f = valid_mask.to(lam_ins.dtype)  # (B, L) 1.0 for valid, 0.0 for pad
+
+        # total outgoing rate at each position
+        total_rate_pos = lam_ins + lam_del + lam_sub     # (B, L)
+        total_rate_pos = total_rate_pos * valid_f        # zero out pads
+        loss_rate = total_rate_pos.sum(dim=1)            # (B,)
+
+        # 2. ----- EDIT TERM -----
+        # precompute log-softmax for tokens (better numerics)
+        logp_ins = F.log_softmax(logits_ins, dim=-1)   # (B, L, V)
+        logp_sub = F.log_softmax(logits_sub, dim=-1)   # (B, L, V)
+
+        # Make weight shape nice
+        if precomputed_weight.dim() == 0:
+            precomputed_weight = precomputed_weight.view(1).expand(B).to(device)
+        else:
+            precomputed_weight = precomputed_weight.to(device)
+
+        # use float for accumulation
+        loss_edit = torch.zeros(B, dtype=torch.float32, device=device)
+
+        for b in range(B):
+            # how many *valid* tokens in x_t[b]
+            valid_len = int(valid_mask[b].sum().item())
+
+            prefix_non_eps = 0  # = number of non-ε seen in z_t[b, :i]
+
+            for i in range(N):
+                zt = int(z_t[b, i].item())
+                z1 = int(z_1[b, i].item())
+
+                # map this column to an x_t position
+                if zt != eps_id:
+                    # this column corresponds to x_t[b, prefix_non_eps]
+                    x_pos = prefix_non_eps
+                    is_token = True
+                    prefix_non_eps += 1
+                else:
+                    # this is a gap column, sits BETWEEN tokens
+                    x_pos = prefix_non_eps
+                    is_token = False
+
+                # if already matched -> no term
+                if zt == z1:
+                    continue
+
+                # --- BOS/EOS guards on the target side ---
+                # if the target token is BOS or EOS, we shouldn't try to force an edit here
+                if z1 == bos_id or z1 == eos_id:
+                    # skip this column; aligned BOS/EOS should already match
+                    continue
+
+                # figure out which edit we need
+                # CASE 1: deletion: z_t has token, z_1 has ε
+                if is_token and (z1 == eps_id):
+                    # delete token at x_pos
+                    if x_pos >= valid_len:
+                        # out of range (shouldn't happen if alignment & mask match)
+                        raise NotImplementedError
+
+                    x_token = int(x_t[b, x_pos].item())
+                    # do NOT delete BOS/EOS
+                    if x_token == bos_id or x_token == eos_id:
+                        continue
+
+                    lam = lam_del[b, x_pos].clamp_min(1e-12)
+                    log_u_req = torch.log(lam)
+
+                # CASE 2: substitution: token -> different token
+                elif is_token and (z1 != eps_id) and (zt != z1):
+                    if x_pos >= valid_len:
+                        raise NotImplementedError
+
+                    x_token = int(x_t[b, x_pos].item())
+                    # do NOT substitute BOS/EOS
+                    if x_token == bos_id or x_token == eos_id:
+                        continue
+
+                    lam = lam_sub[b, x_pos].clamp_min(1e-12)
+                    logp_tok = logp_sub[b, x_pos, z1]
+                    log_u_req = torch.log(lam) + logp_tok
+
+                # CASE 3: insertion: ε -> token
+                elif (not is_token) and (z1 != eps_id):
+                    # insertion in the gap after token (x_pos - 1)
+                    # if x_pos == 0, we insert after "BOS"/at start -> map to position 0
+                    ins_pos = x_pos - 1
+                    if ins_pos < 0:
+                        ins_pos = 0
+
+                    # clamp to last valid position if needed
+                    if valid_len == 0:
+                        # degenerate, but avoid -1
+                        ins_pos = 0
+                    elif ins_pos >= valid_len:
+                        ins_pos = valid_len - 1
+
+                    # also don't insert "after" EOS if ins_pos currently points to EOS
+                    x_token = int(x_t[b, ins_pos].item())
+                    if x_token == eos_id:
+                        # simplest policy: skip this insertion supervision
+                        continue
+
+                    lam = lam_ins[b, ins_pos].clamp_min(1e-12)
+                    logp_tok = logp_ins[b, ins_pos, z1]
+                    log_u_req = torch.log(lam) + logp_tok
+
+                else:
+                    # unknown pattern (shouldn't happen with proper alignment)
+                    raise NotImplementedError
+
+                w = precomputed_weight[b]
+                loss_edit[b] += - w * log_u_req
+
+        # 3. ----- COMBINE -----
+        loss = loss_rate + loss_edit  # (B,)
+        loss = loss.mean()
+        return loss
+
+
+class EditFlowsLossReParam(_Loss):
+    """
+    Edit Flows loss (Eq. 23) but parameterized by:
+      - lam_total: (B, L) >= 0
+      - pi_type:   (B, L, 3) prob over {ins, del, sub}, sums to 1
+
+    Token distributions are still given by logits_ins/logits_sub.
+
+    Required-edit log rate:
+      del: log u = log lam_total[pos] + log pi_del[pos]
+      sub: log u = log lam_total[pos] + log pi_sub[pos] + log Q_sub(token)
+      ins: log u = log lam_total[pos] + log pi_ins[pos] + log Q_ins(token)
+    """
+
+    def __init__(self, reduction: str = "mean", gamma_rate: float = 1.0, eps: float = 1e-12) -> None:
+        super().__init__(None, None, reduction)
+        self.gamma_rate = float(gamma_rate)
+        self.eps = float(eps)
+
+    def forward(
+        self,
+        lam_total: torch.Tensor,     # (B, L)
+        pi_type: torch.Tensor,       # (B, L, 3) probabilities over {ins, del, sub}
+        logits_ins: torch.Tensor,    # (B, L, V)
+        logits_sub: torch.Tensor,    # (B, L, V)
+        z_t: torch.Tensor,           # (B, N) aligned, with eps_id
+        z_1: torch.Tensor,           # (B, N) aligned target, with eps_id
+        x_t: torch.Tensor,           # (B, L)
+        valid_mask: torch.Tensor,    # (B, L) bool, False==padding, True==valid   (same as your old loss)
+        precomputed_weight: torch.Tensor,  # (B,) or ()  = kappa_dot/(1-kappa)
+        eps_id: int,
+        bos_id: int,
+        eos_id: int,
+    ) -> torch.Tensor:
+        device = lam_total.device
+        B, L = x_t.shape
+        _, N = z_t.shape
+
+        # --- RATE TERM ---
+        valid_f = valid_mask.to(lam_total.dtype)                 # True=valid -> 1.0
+        loss_rate = (lam_total * valid_f).sum(dim=1)             # (B,)
+        loss_rate = self.gamma_rate * loss_rate
+
+        # --- EDIT TERM ---
+        logp_ins = F.log_softmax(logits_ins, dim=-1)             # (B, L, V)
+        logp_sub = F.log_softmax(logits_sub, dim=-1)             # (B, L, V)
+
+        # normalize pi_type defensively (in case caller passes slightly off-simplex)
+        pi_sum = pi_type.sum(dim=-1, keepdim=True).clamp_min(self.eps)
+        pi_type = pi_type / pi_sum
+
+        log_lam = torch.log(lam_total.clamp_min(self.eps))       # (B, L)
+        log_pi = torch.log(pi_type.clamp_min(self.eps))          # (B, L, 3)
+
+        # weight shape
+        if precomputed_weight.dim() == 0:
+            precomputed_weight = precomputed_weight.view(1).expand(B).to(device)
+        else:
+            precomputed_weight = precomputed_weight.to(device)
+
+        loss_edit = torch.zeros(B, dtype=torch.float32, device=device)
+
+        for b in range(B):
+            valid_len = int(valid_mask[b].sum().item())
+            prefix_non_eps = 0
+
+            for i in range(N):
+                zt = int(z_t[b, i].item())
+                z1 = int(z_1[b, i].item())
+
+                # map this alignment column to an x_t position
+                if zt != eps_id:
+                    x_pos = prefix_non_eps
+                    is_token = True
+                    prefix_non_eps += 1
+                else:
+                    x_pos = prefix_non_eps
+                    is_token = False
+
+                # already matched -> no supervision
+                if zt == z1:
+                    continue
+
+                # don't force edits where target is BOS/EOS
+                if z1 == bos_id or z1 == eos_id:
+                    continue
+
+                # CASE 1: deletion (token -> eps)
+                if is_token and (z1 == eps_id):
+                    if x_pos >= valid_len:
+                        raise NotImplementedError
+
+                    x_token = int(x_t[b, x_pos].item())
+                    if x_token == bos_id or x_token == eos_id:
+                        continue
+
+                    log_u_req = log_lam[b, x_pos] + log_pi[b, x_pos, 1]  # del index = 1
+
+                # CASE 2: substitution (token -> different token)
+                elif is_token and (z1 != eps_id) and (zt != z1):
+                    if x_pos >= valid_len:
+                        raise NotImplementedError
+
+                    x_token = int(x_t[b, x_pos].item())
+                    if x_token == bos_id or x_token == eos_id:
+                        continue
+
+                    log_u_req = (
+                        log_lam[b, x_pos]
+                        + log_pi[b, x_pos, 2]           # sub index = 2
+                        + logp_sub[b, x_pos, z1]
+                    )
+
+                # CASE 3: insertion (eps -> token), map gap to "after token (x_pos-1)"
+                elif (not is_token) and (z1 != eps_id):
+                    ins_pos = x_pos - 1
+                    if ins_pos < 0:
+                        ins_pos = 0
+
+                    if valid_len == 0:
+                        ins_pos = 0
+                    elif ins_pos >= valid_len:
+                        ins_pos = valid_len - 1
+
+                    x_token = int(x_t[b, ins_pos].item())
+                    if x_token == eos_id:
+                        continue
+
+                    log_u_req = (
+                        log_lam[b, ins_pos]
+                        + log_pi[b, ins_pos, 0]         # ins index = 0
+                        + logp_ins[b, ins_pos, z1]
+                    )
+
+                else:
+                    raise NotImplementedError
+
+                w = precomputed_weight[b]
+                loss_edit[b] += -w * log_u_req
+
+        loss = loss_rate + loss_edit
+        return loss.mean()

flow_matching/loss/generalized_loss.py

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+from torch.nn.modules.loss import _Loss
+
+from flow_matching.path import MixtureDiscreteProbPath
+
+
+class MixturePathGeneralizedKL(_Loss):
+    r"""A generalized KL loss for discrete flow matching.
+    A class that measures the generalized KL of a discrete flow model :math:`p_{1|t}` w.r.t. a probability path given by ``path``. Note: this class is assuming that the model is trained on the same path.
+
+    For a model trained on a space :math:`\mathcal{S} = \mathcal{T}^d`, :math:`\mathcal{T} = [K] = \set{1,2,\ldots,K}`, the loss is given by
+
+    .. math::
+            \ell_i(x_1, x_t, t) = -\frac{\dot{\kappa}_t}{1-\kappa_t} \biggr[  p_{1|t}(x_t^i|x_t) -\delta_{x^i_1}(x_t^i) + (1-\delta_{x^i_1}(x_t^i))\left(\log p_{1|t}(x_1^i|x_t)\right)\biggr],
+
+    where :math:`\kappa_t` is the scheduler associated with ``path``.
+
+    Args:
+        path (MixtureDiscreteProbPath): Probability path (x-prediction training).
+        reduction (str, optional): Specify the reduction to apply to the output ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction is applied to the output, ``'mean'``: the output is reduced by mean over sequence elements, ``'sum'``: the output is reduced by sum over sequence elements. Defaults to 'mean'.
+    """
+
+    def __init__(self, path: MixtureDiscreteProbPath, reduction: str = "mean") -> None:
+        super().__init__(None, None, reduction)
+        self.path = path
+
+    def forward(self, logits: Tensor, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Evaluates the generalized KL loss.
+
+        Args:
+            logits (Tensor): posterior model output (i.e., softmax(``logits``) :math:`=p_{1|t}(x|x_t)`), shape (batch, d, K).
+            x_1 (Tensor): target data point :math:`x_1 \sim q`, shape (batch, d).
+            x_t (Tensor): conditional sample at :math:`x_t \sim p_t(\cdot|x_1)`, shape (batch, d).
+            t (Tensor): times in :math:`[0,1]`, shape (batch).
+
+        Raises:
+            ValueError: reduction value must be one of ``'none'`` | ``'mean'`` | ``'sum'``.
+
+        Returns:
+            Tensor: Generalized KL loss.
+        """
+        x_1_shape = x_1.shape
+
+        # extract x_1 value of log(p_{1|t}(x|x_t)).
+        log_p_1t = torch.log_softmax(logits, dim=-1)
+        log_p_1t_x1 = torch.gather(log_p_1t, dim=-1, index=x_1.unsqueeze(-1))
+        log_p_1t_x1 = log_p_1t_x1.view(*x_1_shape)
+
+        # extract x_t value of p_{1|t}(x|x_t).
+        p_1t = torch.exp(log_p_1t)
+        p_1t_xt = torch.gather(p_1t, dim=-1, index=x_t.unsqueeze(-1))
+        p_1t_xt = p_1t_xt.view(*x_1_shape)
+
+        scheduler_output = self.path.scheduler(t)
+
+        jump_coefficient = (
+            scheduler_output.d_alpha_t / (1 - scheduler_output.alpha_t)
+        )[(...,) + (None,) * (x_1.dim() - 1)]
+        jump_coefficient = jump_coefficient.repeat(1, *x_1_shape[1:])
+        delta_x1_xt = (x_t == x_1).to(log_p_1t.dtype)
+
+        loss = -jump_coefficient * (
+            p_1t_xt - delta_x1_xt + (1 - delta_x1_xt) * log_p_1t_x1
+        )
+
+        if self.reduction == "mean":
+            return torch.mean(loss)
+        elif self.reduction == "sum":
+            return torch.sum(loss)
+        elif self.reduction == "none":
+            return loss
+        else:
+            raise ValueError(f"{self.reduction} is not a valid value for reduction")

flow_matching/path/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .affine import AffineProbPath, CondOTProbPath
+from .geodesic import GeodesicProbPath
+from .mixture import MixtureDiscreteProbPath
+from .path import ProbPath
+from .path_sample import DiscretePathSample, PathSample
+from .editflows_adapter import EditFlowsPathAdapter
+
+
+__all__ = [
+    "ProbPath",
+    "AffineProbPath",
+    "CondOTProbPath",
+    "MixtureDiscreteProbPath",
+    "GeodesicProbPath",
+    "PathSample",
+    "DiscretePathSample",
+    "EditFlowsPathAdapter",
+]

flow_matching/path/affine.py

@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import Tensor
+
+from flow_matching.path.path import ProbPath
+from flow_matching.path.path_sample import PathSample
+from flow_matching.path.scheduler.scheduler import CondOTScheduler, Scheduler
+from flow_matching.utils import expand_tensor_like
+
+
+class AffineProbPath(ProbPath):
+    r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
+    An affine transformation can be represented as:
+
+    .. math::
+
+        X_t = \alpha_t X_1 + \sigma_t X_0,
+
+    where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.
+
+    The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.
+
+    Using ``AffineProbPath`` in the flow matching framework:
+
+    .. code-block:: python
+
+        # Instantiates a probability path
+        my_path = AffineProbPath(...)
+        mse_loss = torch.nn.MSELoss()
+
+        for x_1 in dataset:
+            # Sets x_0 to random noise
+            x_0 = torch.randn()
+
+            # Sets t to a random value in [0,1]
+            t = torch.rand()
+
+            # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+            # Computes the MSE loss w.r.t. the velocity
+            loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+            loss.backward()
+
+    Args:
+        scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.
+
+    """
+
+    def __init__(self, scheduler: Scheduler):
+        self.scheduler = scheduler
+
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+        r"""Sample from the affine probability path:
+
+        | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+        | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.
+
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            PathSample: a conditional sample at :math:`X_t \sim p_t`.
+        """
+        self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.alpha_t, expand_to=x_1
+        )
+        sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.sigma_t, expand_to=x_1
+        )
+        d_alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
+        )
+        d_sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
+        )
+
+        # construct xt ~ p_t(x|x1).
+        x_t = sigma_t * x_0 + alpha_t * x_1
+        dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
+
+        return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)
+
+    def target_to_velocity(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from x_1 representation to velocity.
+
+        | given :math:`X_1`.
+        | return :math:`\dot{X}_t`.
+
+        Args:
+            x_1 (Tensor): target data point.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: velocity.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        d_alpha_t = scheduler_output.d_alpha_t
+        sigma_t = scheduler_output.sigma_t
+        d_sigma_t = scheduler_output.d_sigma_t
+
+        a_t = d_sigma_t / sigma_t
+        b_t = (d_alpha_t * sigma_t - d_sigma_t * alpha_t) / sigma_t
+
+        return a_t * x_t + b_t * x_1
+
+    def epsilon_to_velocity(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from epsilon representation to velocity.
+
+        | given :math:`\epsilon`.
+        | return :math:`\dot{X}_t`.
+
+        Args:
+            epsilon (Tensor): noise in the path sample.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: velocity.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        d_alpha_t = scheduler_output.d_alpha_t
+        sigma_t = scheduler_output.sigma_t
+        d_sigma_t = scheduler_output.d_sigma_t
+
+        a_t = d_alpha_t / alpha_t
+        b_t = (d_sigma_t * alpha_t - d_alpha_t * sigma_t) / alpha_t
+
+        return a_t * x_t + b_t * epsilon
+
+    def velocity_to_target(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from velocity to x_1 representation.
+
+        | given :math:`\dot{X}_t`.
+        | return :math:`X_1`.
+
+        Args:
+            velocity (Tensor): velocity at the path sample.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: target data point.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        d_alpha_t = scheduler_output.d_alpha_t
+        sigma_t = scheduler_output.sigma_t
+        d_sigma_t = scheduler_output.d_sigma_t
+
+        a_t = -d_sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t)
+        b_t = sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t)
+
+        return a_t * x_t + b_t * velocity
+
+    def epsilon_to_target(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from epsilon representation to x_1 representation.
+
+        | given :math:`\epsilon`.
+        | return :math:`X_1`.
+
+        Args:
+            epsilon (Tensor): noise in the path sample.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: target data point.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        sigma_t = scheduler_output.sigma_t
+
+        a_t = 1 / alpha_t
+        b_t = -sigma_t / alpha_t
+
+        return a_t * x_t + b_t * epsilon
+
+    def velocity_to_epsilon(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from velocity to noise representation.
+
+        | given :math:`\dot{X}_t`.
+        | return :math:`\epsilon`.
+
+        Args:
+            velocity (Tensor): velocity at the path sample.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: noise in the path sample.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        d_alpha_t = scheduler_output.d_alpha_t
+        sigma_t = scheduler_output.sigma_t
+        d_sigma_t = scheduler_output.d_sigma_t
+
+        a_t = -d_alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t)
+        b_t = alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t)
+
+        return a_t * x_t + b_t * velocity
+
+    def target_to_epsilon(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+        r"""Convert from x_1 representation to velocity.
+
+        | given :math:`X_1`.
+        | return :math:`\epsilon`.
+
+        Args:
+            x_1 (Tensor): target data point.
+            x_t (Tensor): path sample at time t.
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: noise in the path sample.
+        """
+        scheduler_output = self.scheduler(t)
+
+        alpha_t = scheduler_output.alpha_t
+        sigma_t = scheduler_output.sigma_t
+
+        a_t = 1 / sigma_t
+        b_t = -alpha_t / sigma_t
+
+        return a_t * x_t + b_t * x_1
+
+
+class CondOTProbPath(AffineProbPath):
+    r"""The ``CondOTProbPath`` class represents a conditional optimal transport probability path.
+
+    This class is a specialized version of the ``AffineProbPath`` that uses a conditional optimal transport scheduler to determine the parameters of the affine transformation.
+
+    The parameters :math:`\alpha_t` and :math:`\sigma_t` for the conditional optimal transport path are defined as:
+
+    .. math::
+
+        \alpha_t = t \quad \text{and} \quad \sigma_t = 1 - t.
+    """
+
+    def __init__(self):
+        self.scheduler = CondOTScheduler()

flow_matching/path/editflows_adapter.py

@@ -0,0 +1,85 @@
+from dataclasses import dataclass
+from typing import List, Sequence, Optional, Tuple
+
+import torch
+from torch import Tensor
+import pdb
+
+from flow_matching.path import MixtureDiscreteProbPath
+
+
+@dataclass
+class EditFlowsPathSampleRagged:
+    # aligned inputs/outputs (ragged)
+    z0_list: List[Tensor]      # list[(T_i,)]
+    z1_list: List[Tensor]      # list[(T_i,)]
+    zt_list: List[Tensor]      # list[(T_i,)]
+    # projected current sequences (ragged X_t)
+    x_t_list: List[Tensor]     # list[(n_i,)]
+    # targets/masks for EF loss (ragged)
+    need_delete_list: List[Tensor]        # list[(n_i,)]   bool
+    need_substitute_list: List[Tensor]    # list[(n_i,)]   bool
+    sub_target_list: List[Tensor]         # list[(n_i,)]   long
+    ins_slot_idx_list: List[Tensor]       # list[(K_i,)]   long  indices in [0..n_i]
+    ins_target_list: List[Tensor]         # list[(K_i,)]   long
+    # time
+    t: Tensor                              # (B,)
+
+
+class EditFlowsPathAdapter:
+    """
+    Ragged adapter for Edit Flows:
+      - Accepts aligned Z sequences as tensors (B, N) (may include eps_id).
+      - Applies MixtureDiscreteProbPath's scheduler (σ_t) to sample z_t by
+        flipping z0→z1 with prob (1-σ_t) per aligned token.
+      - You can later project z_t -> x_t by removing ε.
+    """
+    def __init__(self, mixture_path, eps_id: int):
+        self.path = mixture_path
+        self.eps_id = int(eps_id)
+
+    # expose scheduler so training can precompute weights κ̇/(1-κ)
+    def scheduler(self, t: Tensor):
+        return self.path.scheduler(t)
+
+    @torch.no_grad()
+    def sample(
+        self,
+        z0: Tensor,   # (B, N)
+        z1: Tensor,   # (B, N)
+        t: Tensor,    # () or (B,)
+    ) -> Tensor:
+        """
+        Sample z_t from (z0, z1) using the discrete mixture path:
+          - σ_t = scheduler(t)
+          - with prob σ_t keep z0
+          - with prob (1 - σ_t) take z1
+        We do this per-column.
+        BOS/EOS are already aligned -> z0 == z1 there -> safe.
+        """
+        device = z0.device
+        B, N = z0.shape
+
+        # get sigma_t from the path
+        sigma = self.path.scheduler(t).sigma_t  # could be scalar or (B,)
+        if sigma.dim() == 0:
+            sigma = sigma.expand(B)          # (B,)
+        else:
+            # ensure shape is (B,)
+            sigma = sigma.view(B)
+        # reshape to broadcast over columns
+        sigma = sigma.view(B, 1)             # (B, 1)
+
+        # uniform noise per position
+        u = torch.rand(B, N, device=device)
+
+        # we only need to flip where z0 != z1
+        diff_mask = (z0 != z1)               # (B, N)
+
+        # flip where u > sigma (i.e. prob 1 - sigma) AND tokens differ
+        flip = (u > sigma) & diff_mask       # (B, N)
+
+        # choose from z1 when flip, else z0
+        z_t = torch.where(flip, z1, z0)      # (B, N)
+
+        return z_t

flow_matching/path/geodesic.py

@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from torch import Tensor
+from torch.func import jvp, vmap
+
+from flow_matching.path.path import ProbPath
+
+from flow_matching.path.path_sample import PathSample
+from flow_matching.path.scheduler import ConvexScheduler
+from flow_matching.utils import expand_tensor_like
+
+from flow_matching.utils.manifolds import geodesic, Manifold
+
+
+class GeodesicProbPath(ProbPath):
+    r"""The ``GeodesicProbPath`` class represents a specific type of probability path where the transformation between distributions is defined through the geodesic path.
+    Mathematically, a geodesic path can be represented as:
+
+    .. math::
+
+        X_t = \psi_t(X_0 | X_1) = \exp_{X_1}(\kappa_t \log_{X_1}(X_0)),
+
+    where :math:`X_t` is the transformed data point at time `t`, :math:`X_0` and :math:`X_1` are the source and target data points, respectively, and :math:`\kappa_t` is a scheduler.
+
+    The scheduler is responsible for providing the time-dependent :math:`\kappa_t` and must be differentiable.
+
+    Using ``GeodesicProbPath`` in the flow matching framework:
+
+    .. code-block:: python
+        # Instantiates a manifold
+        manifold = FlatTorus()
+
+        # Instantiates a scheduler
+        scheduler = CondOTScheduler()
+
+        # Instantiates a probability path
+        my_path = GeodesicProbPath(scheduler, manifold)
+        mse_loss = torch.nn.MSELoss()
+
+        for x_1 in dataset:
+            # Sets x_0 to random noise
+            x_0 = torch.randn()
+
+            # Sets t to a random value in [0,1]
+            t = torch.rand()
+
+            # Samples the conditional path :math:`X_t \sim p_t(X_t|X_0,X_1)`
+            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+            # Computes the MSE loss w.r.t. the velocity
+            loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+            loss.backward()
+
+    Args:
+        scheduler (ConvexScheduler): The scheduler that provides :math:`\kappa_t`.
+        manifold (Manifold): The manifold on which the probability path is defined.
+
+    """
+
+    def __init__(self, scheduler: ConvexScheduler, manifold: Manifold):
+        self.scheduler = scheduler
+        self.manifold = manifold
+
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+        r"""Sample from the Riemannian probability path with geodesic interpolation:
+
+        | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`\kappa_t`.
+        | return :math:`X_0, X_1, X_t = \exp_{X_1}(\kappa_t \log_{X_1}(X_0))`, and the conditional velocity at :math:`X_t, \dot{X}_t`.
+
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            PathSample: A conditional sample at :math:`X_t \sim p_t`.
+        """
+        self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+        t = expand_tensor_like(input_tensor=t, expand_to=x_1[..., 0:1]).clone()
+
+        def cond_u(x_0, x_1, t):
+            path = geodesic(self.manifold, x_0, x_1)
+            x_t, dx_t = jvp(
+                lambda t: path(self.scheduler(t).alpha_t),
+                (t,),
+                (torch.ones_like(t).to(t),),
+            )
+            return x_t, dx_t
+
+        x_t, dx_t = vmap(cond_u)(x_0, x_1, t)
+        x_t = x_t.reshape_as(x_1)
+        dx_t = dx_t.reshape_as(x_1)
+
+        return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)

flow_matching/path/mixture.py

@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+
+from flow_matching.path.path import ProbPath
+
+from flow_matching.path.path_sample import DiscretePathSample
+from flow_matching.path.scheduler import ConvexScheduler
+from flow_matching.utils import expand_tensor_like, unsqueeze_to_match
+
+
+class MixtureDiscreteProbPath(ProbPath):
+    r"""The ``MixtureDiscreteProbPath`` class defines a factorized discrete probability path.
+
+    This path remains constant at the source data point :math:`X_0` until a random time, determined by the scheduler, when it flips to the target data point :math:`X_1`.
+    The scheduler determines the flip probability using the parameter :math:`\sigma_t`, which is a function of time `t`. Specifically, :math:`\sigma_t` represents the probability of remaining at :math:`X_0`, while :math:`1 - \sigma_t` is the probability of flipping to :math:`X_1`:
+
+    .. math::
+
+        P(X_t = X_0) = \sigma_t \quad \text{and} \quad  P(X_t = X_1) = 1 - \sigma_t,
+
+    where :math:`\sigma_t` is provided by the scheduler.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> x_0 = torch.zeros((1, 3, 3))
+        >>> x_1 = torch.ones((1, 3, 3))
+
+        >>> path = MixtureDiscreteProbPath(PolynomialConvexScheduler(n=1.0))
+        >>> result = path.sample(x_0, x_1, t=torch.tensor([0.1])).x_t
+        >>> result
+        tensor([[[0.0, 0.0, 0.0],
+                 [0.0, 0.0, 1.0],
+                 [0.0, 0.0, 0.0]]])
+
+        >>> result = path.sample(x_0, x_1, t=torch.tensor([0.5])).x_t
+        >>> result
+        tensor([[[1.0, 0.0, 1.0],
+                 [0.0, 1.0, 0.0],
+                 [0.0, 1.0, 0.0]]])
+
+        >>> result = path.sample(x_0, x_1, t=torch.tensor([1.0])).x_t
+        >>> result
+        tensor([[[1.0, 1.0, 1.0],
+                 [1.0, 1.0, 1.0],
+                 [1.0, 1.0, 1.0]]])
+
+    Args:
+        scheduler (ConvexScheduler): The scheduler that provides :math:`\sigma_t`.
+    """
+
+    def __init__(self, scheduler: ConvexScheduler):
+        assert isinstance(
+            scheduler, ConvexScheduler
+        ), "Scheduler for ConvexProbPath must be a ConvexScheduler."
+
+        self.scheduler = scheduler
+
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> DiscretePathSample:
+        r"""Sample from the affine probability path:
+            | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+            | return :math:`X_0, X_1, t`, and :math:`X_t \sim p_t`.
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            DiscretePathSample: a conditional sample at :math:`X_t ~ p_t`.
+        """
+        self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+
+        sigma_t = self.scheduler(t).sigma_t
+
+        sigma_t = expand_tensor_like(input_tensor=sigma_t, expand_to=x_1)
+
+        source_indices = torch.rand(size=x_1.shape, device=x_1.device) < sigma_t
+        x_t = torch.where(condition=source_indices, input=x_0, other=x_1)
+
+        return DiscretePathSample(x_t=x_t, x_1=x_1, x_0=x_0, t=t)
+
+    def posterior_to_velocity(
+        self, posterior_logits: Tensor, x_t: Tensor, t: Tensor
+    ) -> Tensor:
+        r"""Convert the factorized posterior to velocity.
+
+        | given :math:`p(X_1|X_t)`. In the factorized case: :math:`\prod_i p(X_1^i | X_t)`.
+        | return :math:`u_t`.
+
+        Args:
+            posterior_logits (Tensor): logits of the x_1 posterior conditional on x_t, shape (..., vocab size).
+            x_t (Tensor): path sample at time t, shape (...).
+            t (Tensor): time in [0,1].
+
+        Returns:
+            Tensor: velocity.
+        """
+        posterior = torch.softmax(posterior_logits, dim=-1)
+        vocabulary_size = posterior.shape[-1]
+        x_t = F.one_hot(x_t, num_classes=vocabulary_size)
+        t = unsqueeze_to_match(source=t, target=x_t)
+
+        scheduler_output = self.scheduler(t)
+
+        kappa_t = scheduler_output.alpha_t
+        d_kappa_t = scheduler_output.d_alpha_t
+
+        return (d_kappa_t / (1 - kappa_t)) * (posterior - x_t)

flow_matching/path/path.py

@@ -0,0 +1,64 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+
+from torch import Tensor
+
+from flow_matching.path.path_sample import PathSample
+
+
+class ProbPath(ABC):
+    r"""Abstract class, representing a probability path.
+
+    A probability path transforms the distribution :math:`p(X_0)` into :math:`p(X_1)` over :math:`t=0\rightarrow 1`.
+
+    The ``ProbPath`` class is designed to support model training in the flow matching framework. It supports two key functionalities: (1) sampling the conditional probability path and (2) conversion between various training objectives.
+    Here is a high-level example
+
+    .. code-block:: python
+
+        # Instantiate a probability path
+        my_path = ProbPath(...)
+
+        for x_0, x_1 in dataset:
+            # Sets t to a random value in [0,1]
+            t = torch.rand()
+
+            # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+            # Optimizes the model. The loss function varies, depending on model and path.
+            loss(path_sample, my_model(x_t, t)).backward()
+
+    """
+
+    @abstractmethod
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+        r"""Sample from an abstract probability path:
+
+        | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)`.
+        | returns :math:`X_0, X_1, X_t \sim p_t(X_t)`, and a conditional target :math:`Y`, all objects are under ``PathSample``.
+
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            PathSample: a conditional sample.
+        """
+
+    def assert_sample_shape(self, x_0, x_1, t):
+        # x_0, x_1: list[Tensor] of variable length (ragged)
+        # t: Tensor of shape (B,)
+        assert isinstance(x_0, (list, tuple)), f"x_0 must be a list or tuple of ragged Tensors. Got {type(x_0)}"
+        assert isinstance(x_1, (list, tuple)), f"x_1 must be a list or tuple of ragged Tensors. Got {type(x_1)}"
+        assert isinstance(t, Tensor), f"t must be a torch.Tensor. Got {type(t)}"
+        batch_size = len(x_0)
+        assert len(x_1) == batch_size, f"x_1 must match x_0 batch size ({batch_size}). Got {len(x_1)}"
+        assert t.ndim == 1, f"The time vector t must have shape [batch_size]. Got {t.shape}."
+        assert t.shape[0] == batch_size, f"Time t dimension must match the batch size [{batch_size}]. Got {t.shape}"

flow_matching/path/path_sample.py

@@ -0,0 +1,53 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+
+from torch import Tensor
+
+
+@dataclass
+class PathSample:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        x_1 (Tensor): the target sample :math:`X_1`.
+        x_0 (Tensor): the source sample :math:`X_0`.
+        t (Tensor): the time sample :math:`t`.
+        x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...).
+        dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...).
+
+    """
+
+    x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+    x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+    t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+    x_t: Tensor = field(
+        metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."}
+    )
+    dx_t: Tensor = field(
+        metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."}
+    )
+
+
+@dataclass
+class DiscretePathSample:
+    r"""
+    Represents a sample of a conditional-flow generated discrete probability path.
+
+    Attributes:
+        x_1 (Tensor): the target sample :math:`X_1`.
+        x_0 (Tensor): the source sample :math:`X_0`.
+        t (Tensor): the time sample  :math:`t`.
+        x_t (Tensor): the sample along the path  :math:`X_t \sim p_t`.
+    """
+
+    x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+    x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+    t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+    x_t: Tensor = field(
+        metadata={"help": "samples X_t ~ p_t(X_t), shape (batch_size, ...)."}
+    )

flow_matching/path/scheduler/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .schedule_transform import ScheduleTransformedModel
+from .scheduler import (
+    CondOTScheduler,
+    ConvexScheduler,
+    CosineScheduler,
+    LinearVPScheduler,
+    PolynomialConvexScheduler,
+    Scheduler,
+    SchedulerOutput,
+    VPScheduler,
+)
+
+__all__ = [
+    "CondOTScheduler",
+    "CosineScheduler",
+    "ConvexScheduler",
+    "PolynomialConvexScheduler",
+    "ScheduleTransformedModel",
+    "Scheduler",
+    "VPScheduler",
+    "LinearVPScheduler",
+    "SchedulerOutput",
+]

flow_matching/path/scheduler/schedule_transform.py

@@ -0,0 +1,148 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import Tensor
+
+from flow_matching.path.scheduler.scheduler import Scheduler
+from flow_matching.utils import ModelWrapper
+
+
+class ScheduleTransformedModel(ModelWrapper):
+    """
+    Change of scheduler for a velocity model.
+
+    This class wraps a given velocity model and transforms its scheduling
+    to a new scheduler function. It modifies the time
+    dynamics of the model according to the new scheduler while maintaining
+    the original model's behavior.
+
+    Example:
+
+    .. code-block:: python
+
+        import torch
+        from flow_matching.path.scheduler import CondOTScheduler, CosineScheduler, ScheduleTransformedModel
+        from flow_matching.solver import ODESolver
+
+        # Initialize the model and schedulers
+        model = ...
+
+        original_scheduler = CondOTScheduler()
+        new_scheduler = CosineScheduler()
+
+        # Create the transformed model
+        transformed_model = ScheduleTransformedModel(
+            velocity_model=model,
+            original_scheduler=original_scheduler,
+            new_scheduler=new_scheduler
+        )
+
+        # Set up the solver
+        solver = ODESolver(velocity_model=transformed_model)
+
+        x_0 = torch.randn([10, 2])  # Example initial condition
+
+        x_1 = solver.sample(
+            time_steps=torch.tensor([0.0, 1.0]),
+            x_init=x_0,
+            step_size=1/1000
+            )[1]
+
+    Args:
+        velocity_model (ModelWrapper): The original velocity model to be transformed.
+        original_scheduler (Scheduler): The scheduler used by the original model. Must implement the snr_inverse function.
+        new_scheduler (Scheduler): The new scheduler to be applied to the model.
+    """
+
+    def __init__(
+        self,
+        velocity_model: ModelWrapper,
+        original_scheduler: Scheduler,
+        new_scheduler: Scheduler,
+    ):
+        super().__init__(model=velocity_model)
+        self.original_scheduler = original_scheduler
+        self.new_scheduler = new_scheduler
+
+        assert hasattr(self.original_scheduler, "snr_inverse") and callable(
+            getattr(self.original_scheduler, "snr_inverse")
+        ), "The original scheduler must have a callable 'snr_inverse' method."
+
+    def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
+        r"""
+        Compute the transformed marginal velocity field for a new scheduler.
+        This method implements a post-training velocity scheduler change for
+        affine conditional flows. It transforms a generating marginal velocity
+        field :math:`u_t(x)` based on an original scheduler to a new marginal velocity
+        field :math:`\bar{u}_r(x)` based on a different scheduler, while maintaining
+        the same data coupling.
+        The transformation is based on the scale-time (ST) transformation
+        between the two conditional flows, defined as:
+
+        .. math::
+
+            \bar{X}_r = s_r X_{t_r},
+
+        where :math:`X_t` and :math:`\bar{X}_r` are defined by their respective schedulers.
+        The ST transformation is computed as:
+
+        .. math::
+
+            t_r = \rho^{-1}(\bar{\rho}(r)) \quad \text{and} \quad  s_r = \frac{\bar{\sigma}_r}{\sigma_{t_r}}.
+
+        Here, :math:`\rho(t)` is the signal-to-noise ratio (SNR) defined as:
+
+        .. math::
+
+            \rho(t) = \frac{\alpha_t}{\sigma_t}.
+
+        :math:`\bar{\rho}(r)` is similarly defined for the new scheduler.
+        The marginal velocity for the new scheduler is then given by:
+
+        .. math::
+
+            \bar{u}_r(x) = \left(\frac{\dot{s}_r}{s_r}\right) x + s_r \dot{t}_r u_{t_r}\left(\frac{x}{s_r}\right).
+
+        Args:
+            x (Tensor): :math:`x_t`, the input tensor.
+            t (Tensor): The time tensor (denoted as :math:`r` above).
+            **extras: Additional arguments for the model.
+        Returns:
+            Tensor: The transformed velocity.
+        """
+        r = t
+
+        r_scheduler_output = self.new_scheduler(t=r)
+
+        alpha_r = r_scheduler_output.alpha_t
+        sigma_r = r_scheduler_output.sigma_t
+        d_alpha_r = r_scheduler_output.d_alpha_t
+        d_sigma_r = r_scheduler_output.d_sigma_t
+
+        t = self.original_scheduler.snr_inverse(alpha_r / sigma_r)
+
+        t_scheduler_output = self.original_scheduler(t=t)
+
+        alpha_t = t_scheduler_output.alpha_t
+        sigma_t = t_scheduler_output.sigma_t
+        d_alpha_t = t_scheduler_output.d_alpha_t
+        d_sigma_t = t_scheduler_output.d_sigma_t
+
+        s_r = sigma_r / sigma_t
+
+        dt_r = (
+            sigma_t
+            * sigma_t
+            * (sigma_r * d_alpha_r - alpha_r * d_sigma_r)
+            / (sigma_r * sigma_r * (sigma_t * d_alpha_t - alpha_t * d_sigma_t))
+        )
+
+        ds_r = (sigma_t * d_sigma_r - sigma_r * d_sigma_t * dt_r) / (sigma_t * sigma_t)
+
+        u_t = self.model(x=x / s_r, t=t, **extras)
+        u_r = ds_r * x / s_r + dt_r * s_r * u_t
+
+        return u_r

flow_matching/path/scheduler/scheduler.py

@@ -0,0 +1,199 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+from typing import Union
+
+import torch
+
+from torch import Tensor
+
+
+@dataclass
+class SchedulerOutput:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        alpha_t (Tensor): :math:`\alpha_t`, shape (...).
+        sigma_t (Tensor): :math:`\sigma_t`, shape (...).
+        d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
+        d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
+
+    """
+
+    alpha_t: Tensor = field(metadata={"help": "alpha_t"})
+    sigma_t: Tensor = field(metadata={"help": "sigma_t"})
+    d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
+    d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
+
+
+class Scheduler(ABC):
+    """Base Scheduler class."""
+
+    @abstractmethod
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        r"""
+        Args:
+            t (Tensor): times in [0,1], shape (...).
+
+        Returns:
+            SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+        """
+        ...
+
+    @abstractmethod
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        r"""
+        Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+        Args:
+            snr (Tensor): The signal-to-noise, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        ...
+
+
+class ConvexScheduler(Scheduler):
+    @abstractmethod
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        r"""Scheduler for convex paths.
+
+        Args:
+            t (Tensor): times in [0,1], shape (...).
+
+        Returns:
+            SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+        """
+        ...
+
+    @abstractmethod
+    def kappa_inverse(self, kappa: Tensor) -> Tensor:
+        r"""
+        Computes :math:`t` from :math:`\kappa_t`.
+
+        Args:
+            kappa (Tensor): :math:`\kappa`, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        ...
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        r"""
+        Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+        Args:
+            snr (Tensor): The signal-to-noise, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        kappa_t = snr / (1.0 + snr)
+
+        return self.kappa_inverse(kappa=kappa_t)
+
+
+class CondOTScheduler(ConvexScheduler):
+    """CondOT Scheduler."""
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        return SchedulerOutput(
+            alpha_t=t,
+            sigma_t=1 - t,
+            d_alpha_t=torch.ones_like(t),
+            d_sigma_t=-torch.ones_like(t),
+        )
+
+    def kappa_inverse(self, kappa: Tensor) -> Tensor:
+        return kappa
+
+
+class PolynomialConvexScheduler(ConvexScheduler):
+    """Polynomial Scheduler."""
+
+    def __init__(self, n: Union[float, int]) -> None:
+        assert isinstance(
+            n, (float, int)
+        ), f"`n` must be a float or int. Got {type(n)=}."
+        assert n > 0, f"`n` must be positive. Got {n=}."
+
+        self.n = n
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        return SchedulerOutput(
+            alpha_t=t**self.n,
+            sigma_t=1 - t**self.n,
+            d_alpha_t=self.n * (t ** (self.n - 1)),
+            d_sigma_t=-self.n * (t ** (self.n - 1)),
+        )
+
+    def kappa_inverse(self, kappa: Tensor) -> Tensor:
+        return torch.pow(kappa, 1.0 / self.n)
+
+
+class VPScheduler(Scheduler):
+    """Variance Preserving Scheduler."""
+
+    def __init__(self, beta_min: float = 0.1, beta_max: float = 20.0) -> None:
+        self.beta_min = beta_min
+        self.beta_max = beta_max
+        super().__init__()
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        b = self.beta_min
+        B = self.beta_max
+        T = 0.5 * (1 - t) ** 2 * (B - b) + (1 - t) * b
+        dT = -(1 - t) * (B - b) - b
+
+        return SchedulerOutput(
+            alpha_t=torch.exp(-0.5 * T),
+            sigma_t=torch.sqrt(1 - torch.exp(-T)),
+            d_alpha_t=-0.5 * dT * torch.exp(-0.5 * T),
+            d_sigma_t=0.5 * dT * torch.exp(-T) / torch.sqrt(1 - torch.exp(-T)),
+        )
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        T = -torch.log(snr**2 / (snr**2 + 1))
+        b = self.beta_min
+        B = self.beta_max
+        t = 1 - ((-b + torch.sqrt(b**2 + 2 * (B - b) * T)) / (B - b))
+        return t
+
+
+class LinearVPScheduler(Scheduler):
+    """Linear Variance Preserving Scheduler."""
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        return SchedulerOutput(
+            alpha_t=t,
+            sigma_t=(1 - t**2) ** 0.5,
+            d_alpha_t=torch.ones_like(t),
+            d_sigma_t=-t / (1 - t**2) ** 0.5,
+        )
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        return torch.sqrt(snr**2 / (1 + snr**2))
+
+
+class CosineScheduler(Scheduler):
+    """Cosine Scheduler."""
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        pi = torch.pi
+        return SchedulerOutput(
+            alpha_t=torch.sin(pi / 2 * t),
+            sigma_t=torch.cos(pi / 2 * t),
+            d_alpha_t=pi / 2 * torch.cos(pi / 2 * t),
+            d_sigma_t=-pi / 2 * torch.sin(pi / 2 * t),
+        )
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        return 2.0 * torch.atan(snr) / torch.pi

flow_matching/solver/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .discrete_solver import MixtureDiscreteEulerSolver
+from .ode_solver import ODESolver
+from .riemannian_ode_solver import RiemannianODESolver
+from .solver import Solver
+from .editflows_solver import EditFlowsEulerSolver
+
+__all__ = [
+    "ODESolver",
+    "Solver",
+    "ModelWrapper",
+    "MixtureDiscreteEulerSolver",
+    "RiemannianODESolver",
+    "EditFlowsEulerSolver",
+]

flow_matching/solver/discrete_solver.py

@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from contextlib import nullcontext
+from math import ceil
+from typing import Callable, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch.nn import functional as F
+
+from flow_matching.path import MixtureDiscreteProbPath
+
+from flow_matching.solver.solver import Solver
+from flow_matching.utils import categorical, ModelWrapper
+from .utils import get_nearest_times
+
+try:
+    from tqdm import tqdm
+
+    TQDM_AVAILABLE = True
+except ImportError:
+    TQDM_AVAILABLE = False
+
+
+class MixtureDiscreteEulerSolver(Solver):
+    r"""Solver that simulates the CTMC process :math:`(X_t)_{t_{\text{init}}\leq t\leq t_{\text{final}}}` defined by :math:`p_t` the marginal probability path of ``path``.
+    Given :math:`X_t \sim p_t`, the algorithm of solver step from :math:`t` to :math:`t+h` for the i-th coordinate is:
+
+    .. math::
+
+        \begin{align*}
+            & X_1^i \sim p_{1|t}^i(\cdot|X_t)\\
+            & \lambda^i \gets \sum_{x^i\ne X_t^i} u_t^i(x^i, X_t^i|X_1^i)\\
+            & Z^i_{\text{change}} \sim U[0,1]\\
+            & X_{t+h}^i \sim \begin{cases}
+                \frac{u_t^i(\cdot, X_t^i|X_1^i)}{\lambda^i}(1-\delta_{X_t^i}(\cdot)) \text{ if $Z^i_{\text{change}}\le 1-e^{-h\lambda^i}$}\\
+                \delta_{X_t^i}(\cdot) \text{ else }
+            \end{cases}
+        \end{align*}
+
+    Where :math:`p_{1|t}(\cdot|X_t)` is the output of ``model``, and the conditional probability velocity is of the mixture probability path is:
+
+    .. math::
+
+        u_t^i(x^i, y^i|x_1^i) = \hat{u}_t^i(x^i, y^i|x_1^i) + c_{\text{div\_free}}\left[\hat{u}_t^i(x^i, y^i|x_1^i) - \check{u}_t^i(x^i, y^i|x_1^i) \right],
+
+    where
+
+    .. math::
+        \hat{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{1-\kappa_t} \left[ \delta_{x_1^i}(x^i) - \delta_{y^i}(x^i) \right],
+
+    and
+
+    .. math::
+
+        \check{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{\kappa_t}\left[ \delta_{y^i}(x^i) - p(x^i) \right].
+
+    The source distribution :math:`p(x^i)` is given by ``p``.
+
+    Args:
+        model (ModelWrapper): trained with x-prediction, outputting posterior probabilities (in the range :math:`[0,1]`), output must be [..., vocabulary_size].
+        path (MixtureDiscreteProbPath): Probability path used for x-prediction training.
+        vocabulary_size (int): size of the discrete vocabulary.
+        source_distribution_p (Optional[Tensor], optional): Source distribution, must be of shape [vocabulary_size]. Required only when divergence-free term for the probability velocity is non-zero. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model: ModelWrapper,
+        path: MixtureDiscreteProbPath,
+        vocabulary_size: int,
+        source_distribution_p: Optional[Tensor] = None,
+    ):
+        super().__init__()
+        self.model = model
+        self.path = path
+        self.vocabulary_size = vocabulary_size
+
+        if source_distribution_p is not None:
+            assert source_distribution_p.shape == torch.Size(
+                [vocabulary_size]
+            ), f"Source distribution p dimension must match the vocabulary size {vocabulary_size}. Got {source_distribution_p.shape}."
+
+        self.source_distribution_p = source_distribution_p
+
+    @torch.no_grad()
+    def sample(
+        self,
+        x_init: Tensor,
+        step_size: Optional[float],
+        div_free: Union[float, Callable[[float], float]] = 0.0,
+        dtype_categorical: torch.dtype = torch.float32,
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        verbose: bool = False,
+        **model_extras,
+    ) -> Tensor:
+        """
+        Sample a sequence of discrete values from the given model.
+
+        .. code-block:: python
+
+            import torch
+            from flow_matching.utils import ModelWrapper
+            from flow_matching.solver import MixtureDiscreteEulerSolver
+
+            class DummyModel(ModelWrapper):
+                def __init__(self):
+                    super().__init__(None)
+                def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
+                    return ...
+
+            model = DummyModel()
+            solver = MixtureDiscreteEulerSolver(model=model)
+
+            x_init = torch.LongTensor([122, 725])
+            step_size = 0.001
+            time_grid = torch.tensor([0.0, 1.0])
+
+            result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
+
+        Args:
+            x_init (Tensor): The initial state.
+            step_size (Optional[float]): If float then time discretization is uniform with the given step size. If None then time discretization is set to be time_grid.
+            div_free (Union[float, Callable[[float], float]]): The coefficient of the divergence-free term in the probability velocity. Can be either a float or a time dependent function. Defaults to 0.0.
+            dtype_categorical (torch.dtype): Precision to use for categorical sampler. Defaults to torch.float32.
+            time_grid (Tensor): The CTMC process is solved in the interval [time_grid[0], time_grid[-1]] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]).
+            return_intermediates (bool): If True then return intermediate time steps according to time_grid. Defaults to False.
+            verbose (bool): Whether to print progress bars. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+            Tensor: The sampled sequence of discrete values.
+
+        Raises:
+            ImportError: To run in verbose mode, tqdm must be installed.
+        """
+        if not div_free == 0.0:
+            assert (
+                self.source_distribution_p is not None
+            ), "Source distribution p must be specified in order to add a divergence-free term to the probability velocity."
+
+        # Initialize the current state `x_t` with the initial state `X_0`.
+        time_grid = time_grid.to(device=x_init.device)
+
+        if step_size is None:
+            # If step_size is None then set the t discretization to time_grid.
+            t_discretization = time_grid
+            n_steps = len(time_grid) - 1
+        else:
+            # If step_size is float then t discretization is uniform with step size set by step_size.
+            t_init = time_grid[0].item()
+            t_final = time_grid[-1].item()
+            assert (
+                t_final - t_init
+            ) > step_size, f"Time interval [time_grid[0], time_grid[-1]] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}."
+
+            n_steps = ceil((t_final - t_init) / step_size)
+            t_discretization = torch.tensor(
+                [t_init + step_size * i for i in range(n_steps)] + [t_final],
+                device=x_init.device,
+            )
+
+            if return_intermediates:
+                # get order of intermediate steps:
+                order = torch.argsort(time_grid)
+                # Compute intermediate steps to return via nearest points in t_discretization to time_grid.
+                time_grid = get_nearest_times(
+                    time_grid=time_grid, t_discretization=t_discretization
+                )
+
+        x_t = x_init.clone()
+        steps_counter = 0
+        res = []
+
+        if return_intermediates:
+            res = [x_init.clone()]
+
+        if verbose:
+            if not TQDM_AVAILABLE:
+                raise ImportError(
+                    "tqdm is required for verbose mode. Please install it."
+                )
+            ctx = tqdm(total=t_final, desc=f"NFE: {steps_counter}")
+        else:
+            ctx = nullcontext()
+
+        with ctx:
+            for i in range(n_steps):
+                t = t_discretization[i : i + 1]
+                h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1]
+
+                # Sample x_1 ~ p_1|t( \cdot |x_t)
+                p_1t = self.model(x=x_t, t=t.repeat(x_t.shape[0]), **model_extras)
+                x_1 = categorical(p_1t.to(dtype=dtype_categorical))
+
+                # Checks if final step
+                if i == n_steps - 1:
+                    x_t = x_1
+                else:
+                    # Compute u_t(x|x_t,x_1)
+                    scheduler_output = self.path.scheduler(t=t)
+
+                    k_t = scheduler_output.alpha_t
+                    d_k_t = scheduler_output.d_alpha_t
+
+                    delta_1 = F.one_hot(x_1, num_classes=self.vocabulary_size).to(
+                        k_t.dtype
+                    )
+                    u = d_k_t / (1 - k_t) * delta_1
+
+                    # Add divergence-free part
+                    div_free_t = div_free(t) if callable(div_free) else div_free
+
+                    if div_free_t > 0:
+                        p_0 = self.source_distribution_p[(None,) * x_t.dim()]
+                        u = u + div_free_t * d_k_t / (k_t * (1 - k_t)) * (
+                            (1 - k_t) * p_0 + k_t * delta_1
+                        )
+
+                    # Set u_t(x_t|x_t,x_1) = 0
+                    delta_t = F.one_hot(x_t, num_classes=self.vocabulary_size)
+                    u = torch.where(
+                        delta_t.to(dtype=torch.bool), torch.zeros_like(u), u
+                    )
+
+                    # Sample x_t ~ u_t( \cdot |x_t,x_1)
+                    intensity = u.sum(dim=-1)  # Assuming u_t(xt|xt,x1) := 0
+                    mask_jump = torch.rand(
+                        size=x_t.shape, device=x_t.device
+                    ) < 1 - torch.exp(-h * intensity)
+
+                    if mask_jump.sum() > 0:
+                        x_t[mask_jump] = categorical(
+                            u[mask_jump].to(dtype=dtype_categorical)
+                        )
+
+                steps_counter += 1
+                t = t + h
+
+                if return_intermediates and (t in time_grid):
+                    res.append(x_t.clone())
+
+                if verbose:
+                    ctx.n = t.item()
+                    ctx.refresh()
+                    ctx.set_description(f"NFE: {steps_counter}")
+
+        if return_intermediates:
+            if step_size is None:
+                return torch.stack(res, dim=0)
+            else:
+                return torch.stack(res, dim=0)[order]
+        else:
+            return x_t

flow_matching/solver/editflows_solver.py

@@ -0,0 +1,168 @@
+# Copyright (c) Meta Platforms, Inc.
+# All rights reserved.
+
+from typing import Callable, List, Sequence, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class EditFlowsEulerSolver:
+    """
+    Ragged EditFlows CTMC solver with Euler thinning (≤1 jump per step).
+
+    At each step h:
+      - query model for (λ_ins[slots], Q_ins[slots], λ_del[tokens], λ_sub[tokens], Q_sub[tokens])
+      - total intensity Λ = sum(λ_ins) + sum(λ_del) + sum(λ_sub)
+      - with prob 1 - exp(-h Λ), take ONE jump:
+           * sample which event from concatenated rates
+           * if ins/sub: sample token from the Q at that position/slot
+           * apply edit to the ragged sequence (insert/delete/replace)
+      - advance t ← t + h
+
+    This is faithful to the CTMC in the small-h regime used by Euler schemes.
+    """
+
+    def __init__(
+        self,
+        model,                                        # callable: (x_list, t_vec) -> ragged heads
+        scheduler: Callable[[Tensor], object],        # path.scheduler(t) providing alpha_t, d_alpha_t, sigma_t if needed
+        vocab_size: int,
+        dtype_categorical: torch.dtype = torch.float64,
+        generator: Optional[torch.Generator] = None,
+    ) -> None:
+        self.model = model
+        self.scheduler = scheduler
+        self.vocab_size = vocab_size
+        self.dtype_cat = dtype_categorical
+        self.gen = generator
+
+    @torch.no_grad()
+    def _sample_categorical_logits(self, logits: Tensor) -> Tensor:
+        # logits: (V,) -> sample id
+        probs = F.softmax(logits, dim=-1, dtype=self.dtype_cat)
+        return torch.multinomial(probs, num_samples=1, replacement=True, generator=self.gen).squeeze(0)
+
+    @torch.no_grad()
+    def _maybe_one_jump(
+        self,
+        x: Tensor,                     # (n,)
+        lam_ins: Tensor,               # (n+1,)
+        logits_ins: Tensor,            # (n+1, V)
+        lam_del: Tensor,               # (n,)
+        lam_sub: Tensor,               # (n,)
+        logits_sub: Tensor,            # (n, V)
+        h: float,
+        device: torch.device,
+    ) -> Tensor:
+        """
+        Apply at most one edit event to x with step size h.
+        Returns possibly modified x (1-D LongTensor).
+        """
+        n = x.numel()
+        # Total intensity
+        sum_ins = lam_ins.sum() if lam_ins.numel() else torch.tensor(0.0, device=device)
+        sum_del = lam_del.sum() if lam_del.numel() else torch.tensor(0.0, device=device)
+        sum_sub = lam_sub.sum() if lam_sub.numel() else torch.tensor(0.0, device=device)
+        total = (sum_ins + sum_del + sum_sub).clamp(min=0.0)
+
+        if total.item() <= 0:
+            return x  # no possible jumps
+
+        # Bernoulli: does a jump occur?
+        p_jump = 1.0 - torch.exp(torch.tensor(-h, device=device) * total)
+        if torch.rand((), device=device, generator=self.gen) >= p_jump:
+            return x  # no jump this step
+
+        # Build event vector: [ins(0..n), del(0..n-1), sub(0..n-1)]
+        parts = []
+        if lam_ins.numel(): parts.append(lam_ins)
+        if lam_del.numel(): parts.append(lam_del)
+        if lam_sub.numel(): parts.append(lam_sub)
+        lam_cat = torch.cat(parts, dim=0)  # (n+1 + n + n,)
+        probs = (lam_cat / total).to(self.dtype_cat)
+
+        # Sample which event
+        e_idx = torch.multinomial(probs, num_samples=1, generator=self.gen).item()
+
+        # Decode which branch
+        offset = 0
+        if e_idx < lam_ins.numel():
+            # INSERT at slot s
+            s = e_idx
+            # sample token from slot distribution
+            y = self._sample_categorical_logits(logits_ins[s])  # id ∈ [0..V-1]
+            # perform insertion at slot s: x' = [x[:s], y, x[s:]]
+            if s == 0:
+                x_new = torch.cat([y.view(1), x], dim=0)
+            elif s == n:
+                x_new = torch.cat([x, y.view(1)], dim=0)
+            else:
+                x_new = torch.cat([x[:s], y.view(1), x[s:]], dim=0)
+            return x_new
+
+        e_idx -= lam_ins.numel()
+        if e_idx < lam_del.numel():
+            # DELETE token j
+            j = e_idx
+            if n == 0:
+                return x
+            if j == 0:
+                return x[1:]
+            elif j == n - 1:
+                return x[:-1]
+            else:
+                return torch.cat([x[:j], x[j + 1 :]], dim=0)
+
+        e_idx -= lam_del.numel()
+        # SUBSTITUTE at token j
+        j = e_idx
+        if n == 0:
+            return x
+        y = self._sample_categorical_logits(logits_sub[j])
+        x_new = x.clone()
+        x_new[j] = y
+        return x_new
+
+    @torch.no_grad()
+    def sample(
+        self,
+        x_list: Sequence[Tensor],   # ragged: list[(n_i,)]
+        n_steps: int,
+        t0: float = 0.0,
+        t1: float = 1.0,
+        verbose: bool = False,
+    ) -> List[Tensor]:
+        """
+        Evolve all sequences independently with shared time grid t_k.
+        """
+        assert n_steps >= 1
+        B = len(x_list)
+        device = x_list[0].device if B > 0 else torch.device("cpu")
+
+        x = [xi.clone() for xi in x_list]
+        t = torch.full((B,), float(t0), device=device)
+        h = (t1 - t0) / float(n_steps)
+
+        for k in range(n_steps):
+            # Query model once at time t_k (vector of size B)
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub = self.model(x, t)
+
+            # Apply at most one jump per sample
+            for i in range(B):
+                x[i] = self._maybe_one_jump(
+                    x=x[i],
+                    lam_ins=lam_ins[i].squeeze(0) if lam_ins[i].dim() == 2 else lam_ins[i],
+                    logits_ins=logits_ins[i].squeeze(0) if logits_ins[i].dim() == 3 else logits_ins[i],
+                    lam_del=lam_del[i].squeeze(0) if lam_del[i].dim() == 2 else lam_del[i],
+                    lam_sub=lam_sub[i].squeeze(0) if lam_sub[i].dim() == 2 else lam_sub[i],
+                    logits_sub=logits_sub[i].squeeze(0) if logits_sub[i].dim() == 3 else logits_sub[i],
+                    h=h,
+                    device=device,
+                )
+
+            # advance time
+            t = t + h
+
+        return x

flow_matching/solver/ode_solver.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import Tensor
+from torchdiffeq import odeint
+
+from flow_matching.solver.solver import Solver
+from flow_matching.utils import gradient, ModelWrapper
+
+
+class ODESolver(Solver):
+    """A class to solve ordinary differential equations (ODEs) using a specified velocity model.
+
+    This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers.
+
+    Args:
+        velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)`
+    """
+
+    def __init__(self, velocity_model: Union[ModelWrapper, Callable]):
+        super().__init__()
+        self.velocity_model = velocity_model
+
+    def sample(
+        self,
+        x_init: Tensor,
+        step_size: Optional[float],
+        method: str = "euler",
+        atol: float = 1e-5,
+        rtol: float = 1e-5,
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        enable_grad: bool = False,
+        **model_extras,
+    ) -> Union[Tensor, Sequence[Tensor]]:
+        r"""Solve the ODE with the velocity field.
+
+        Example:
+
+        .. code-block:: python
+
+            import torch
+            from flow_matching.utils import ModelWrapper
+            from flow_matching.solver import ODESolver
+
+            class DummyModel(ModelWrapper):
+                def __init__(self):
+                    super().__init__(None)
+
+                def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
+                    return torch.ones_like(x) * 3.0 * t**2
+
+            velocity_model = DummyModel()
+            solver = ODESolver(velocity_model=velocity_model)
+            x_init = torch.tensor([0.0, 0.0])
+            step_size = 0.001
+            time_grid = torch.tensor([0.0, 1.0])
+
+            result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
+
+        Args:
+            x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...].
+            step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+            method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+            atol (float): Absolute tolerance, used for adaptive step solvers.
+            rtol (float): Relative tolerance, used for adaptive step solvers.
+            time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]).
+            return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+            enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+            Union[Tensor, Sequence[Tensor]]: The last timestep when return_intermediates=False, otherwise all values specified in time_grid.
+        """
+
+        time_grid = time_grid.to(x_init.device)
+
+        def ode_func(t, x):
+            return self.velocity_model(x=x, t=t, **model_extras)
+
+        ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+        with torch.set_grad_enabled(enable_grad):
+            # Approximate ODE solution with numerical ODE solver
+            sol = odeint(
+                ode_func,
+                x_init,
+                time_grid,
+                method=method,
+                options=ode_opts,
+                atol=atol,
+                rtol=rtol,
+            )
+
+        if return_intermediates:
+            return sol
+        else:
+            return sol[-1]
+
+    def compute_likelihood(
+        self,
+        x_1: Tensor,
+        log_p0: Callable[[Tensor], Tensor],
+        step_size: Optional[float],
+        method: str = "euler",
+        atol: float = 1e-5,
+        rtol: float = 1e-5,
+        time_grid: Tensor = torch.tensor([1.0, 0.0]),
+        return_intermediates: bool = False,
+        exact_divergence: bool = False,
+        enable_grad: bool = False,
+        **model_extras,
+    ) -> Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]:
+        r"""Solve for log likelihood given a target sample at :math:`t=0`.
+
+        Works similarly to sample, but solves the ODE in reverse to compute the log-likelihood. The velocity model must be differentiable with respect to x.
+        The function assumes log_p0 is the log probability of the source distribution at :math:`t=0`.
+
+        Args:
+            x_1 (Tensor): target sample (e.g., samples :math:`X_1 \sim p_1`).
+            log_p0 (Callable[[Tensor], Tensor]): Log probability function of the source distribution.
+            step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+            method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+            atol (float): Absolute tolerance, used for adaptive step solvers.
+            rtol (float): Relative tolerance, used for adaptive step solvers.
+            time_grid (Tensor): If step_size is None then time discretization is set by the time grid. Must start at 1.0 and end at 0.0, otherwise the likelihood computation is not valid. Defaults to torch.tensor([1.0, 0.0]).
+            return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Otherwise only return the final sample. Defaults to False.
+            exact_divergence (bool): Whether to compute the exact divergence or use the Hutchinson estimator.
+            enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+            Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]: Samples at time_grid and log likelihood values of given x_1.
+        """
+        assert (
+            time_grid[0] == 1.0 and time_grid[-1] == 0.0
+        ), f"Time grid must start at 1.0 and end at 0.0. Got {time_grid}"
+
+        # Fix the random projection for the Hutchinson divergence estimator
+        if not exact_divergence:
+            z = (torch.randn_like(x_1).to(x_1.device) < 0) * 2.0 - 1.0
+
+        def ode_func(x, t):
+            return self.velocity_model(x=x, t=t, **model_extras)
+
+        def dynamics_func(t, states):
+            xt = states[0]
+            with torch.set_grad_enabled(True):
+                xt.requires_grad_()
+                ut = ode_func(xt, t)
+
+                if exact_divergence:
+                    # Compute exact divergence
+                    div = 0
+                    for i in range(ut.flatten(1).shape[1]):
+                        div += gradient(ut[:, i], xt, create_graph=True)[:, i].detach()
+                else:
+                    # Compute Hutchinson divergence estimator E[z^T D_x(ut) z]
+                    ut_dot_z = torch.einsum(
+                        "ij,ij->i", ut.flatten(start_dim=1), z.flatten(start_dim=1)
+                    )
+                    grad_ut_dot_z = gradient(ut_dot_z, xt)
+                    div = torch.einsum(
+                        "ij,ij->i",
+                        grad_ut_dot_z.flatten(start_dim=1),
+                        z.flatten(start_dim=1),
+                    )
+
+            return ut.detach(), div.detach()
+
+        y_init = (x_1, torch.zeros(x_1.shape[0], device=x_1.device))
+        ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+        with torch.set_grad_enabled(enable_grad):
+            sol, log_det = odeint(
+                dynamics_func,
+                y_init,
+                time_grid,
+                method=method,
+                options=ode_opts,
+                atol=atol,
+                rtol=rtol,
+            )
+
+        x_source = sol[-1]
+        source_log_p = log_p0(x_source)
+
+        if return_intermediates:
+            return sol, source_log_p + log_det[-1]
+        else:
+            return sol[-1], source_log_p + log_det[-1]

flow_matching/solver/riemannian_ode_solver.py

@@ -0,0 +1,261 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Callable
+
+import torch
+from torch import Tensor
+
+from flow_matching.solver.solver import Solver
+from flow_matching.utils import ModelWrapper
+from flow_matching.utils.manifolds import geodesic, Manifold
+
+try:
+    from tqdm import tqdm
+
+    TQDM_AVAILABLE = True
+except ImportError:
+    TQDM_AVAILABLE = False
+
+
+class RiemannianODESolver(Solver):
+    r"""Riemannian ODE solver
+    Initialize the ``RiemannianODESolver``.
+
+    Args:
+        manifold (Manifold): the manifold to solve on.
+        velocity_model (ModelWrapper): a velocity field model receiving :math:`(x,t)`
+            and returning :math:`u_t(x)` which is assumed to lie on the tangent plane at `x`.
+    """
+
+    def __init__(self, manifold: Manifold, velocity_model: ModelWrapper):
+        super().__init__()
+        self.manifold = manifold
+        self.velocity_model = velocity_model
+
+    def sample(
+        self,
+        x_init: Tensor,
+        step_size: float,
+        projx: bool = True,
+        proju: bool = True,
+        method: str = "euler",
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        verbose: bool = False,
+        enable_grad: bool = False,
+        **model_extras,
+    ) -> Tensor:
+        r"""Solve the ODE with the `velocity_field` on the manifold.
+
+        Args:
+            x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`).
+            step_size (float): The step size.
+            projx (bool): Whether to project the point onto the manifold at each step. Defaults to True.
+            proju (bool): Whether to project the vector field onto the tangent plane at each step. Defaults to True.
+            method (str): One of ["euler", "midpoint", "rk4"]. Defaults to "euler".
+            time_grid (Tensor, optional): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]).
+            return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+            verbose (bool, optional): Whether to print progress bars. Defaults to False.
+            enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+            Tensor: The sampled sequence. Defaults to returning samples at :math:`t=1`.
+
+        Raises:
+            ImportError: To run in verbose mode, tqdm must be installed.
+        """
+        step_fns = {
+            "euler": _euler_step,
+            "midpoint": _midpoint_step,
+            "rk4": _rk4_step,
+        }
+        assert method in step_fns.keys(), f"Unknown method {method}"
+        step_fn = step_fns[method]
+
+        def velocity_func(x, t):
+            return self.velocity_model(x=x, t=t, **model_extras)
+
+        # --- Factor this out.
+        time_grid = torch.sort(time_grid.to(device=x_init.device)).values
+
+        if step_size is None:
+            # If step_size is None then set the t discretization to time_grid.
+            t_discretization = time_grid
+            n_steps = len(time_grid) - 1
+        else:
+            # If step_size is float then t discretization is uniform with step size set by step_size.
+            t_init = time_grid[0].item()
+            t_final = time_grid[-1].item()
+            assert (
+                t_final - t_init
+            ) > step_size, f"Time interval [min(time_grid), max(time_grid)] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}."
+
+            n_steps = math.ceil((t_final - t_init) / step_size)
+            t_discretization = torch.tensor(
+                [step_size * i for i in range(n_steps)] + [t_final],
+                device=x_init.device,
+            )
+        # ---
+        t0s = t_discretization[:-1]
+
+        if verbose:
+            if not TQDM_AVAILABLE:
+                raise ImportError(
+                    "tqdm is required for verbose mode. Please install it."
+                )
+            t0s = tqdm(t0s)
+
+        if return_intermediates:
+            xts = []
+            i_ret = 0
+
+        with torch.set_grad_enabled(enable_grad):
+            xt = x_init
+            for t0, t1 in zip(t0s, t_discretization[1:]):
+                dt = t1 - t0
+                xt_next = step_fn(
+                    velocity_func,
+                    xt,
+                    t0,
+                    dt,
+                    manifold=self.manifold,
+                    projx=projx,
+                    proju=proju,
+                )
+                if return_intermediates:
+                    while (
+                        i_ret < len(time_grid)
+                        and t0 <= time_grid[i_ret]
+                        and time_grid[i_ret] <= t1
+                    ):
+                        xts.append(
+                            interp(self.manifold, xt, xt_next, t0, t1, time_grid[i_ret])
+                        )
+                        i_ret += 1
+                xt = xt_next
+
+        if return_intermediates:
+            return torch.stack(xts, dim=0)
+        else:
+            return xt
+
+
+def interp(manifold, xt, xt_next, t, t_next, t_ret):
+    return geodesic(manifold, xt, xt_next)(
+        (t_ret - t) / (t_next - t).reshape(1)
+    ).reshape_as(xt)
+
+
+def _euler_step(
+    velocity_model: Callable,
+    xt: Tensor,
+    t0: Tensor,
+    dt: Tensor,
+    manifold: Manifold,
+    projx: bool = True,
+    proju: bool = True,
+) -> Tensor:
+    r"""Perform an Euler step on a manifold.
+
+    Args:
+        velocity_model (Callable): the velocity model
+        xt (Tensor): tensor containing the state at time t0
+        t0 (Tensor): the time at which this step is taken
+        dt (Tensor): the step size
+        manifold (Manifold): a manifold object
+        projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+        proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+    Returns:
+        Tensor: tensor containing the state after the step
+    """
+    velocity_fn = lambda x, t: (
+        manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+    )
+    projx_fn = lambda x: manifold.projx(x) if projx else x
+
+    vt = velocity_fn(xt, t0)
+
+    xt = xt + dt * vt
+
+    return projx_fn(xt)
+
+
+def _midpoint_step(
+    velocity_model: Callable,
+    xt: Tensor,
+    t0: Tensor,
+    dt: Tensor,
+    manifold: Manifold,
+    projx: bool = True,
+    proju: bool = True,
+) -> Tensor:
+    r"""Perform a midpoint step on a manifold.
+
+    Args:
+        velocity_model (Callable): the velocity model
+        xt (Tensor): tensor containing the state at time t0
+        t0 (Tensor): the time at which this step is taken
+        dt (Tensor): the step size
+        manifold (Manifold): a manifold object
+        projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+        proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+    Returns:
+        Tensor: tensor containing the state after the step
+    """
+    velocity_fn = lambda x, t: (
+        manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+    )
+    projx_fn = lambda x: manifold.projx(x) if projx else x
+
+    half_dt = 0.5 * dt
+    vt = velocity_fn(xt, t0)
+    x_mid = xt + half_dt * vt
+    x_mid = projx_fn(x_mid)
+
+    xt = xt + dt * velocity_fn(x_mid, t0 + half_dt)
+
+    return projx_fn(xt)
+
+
+def _rk4_step(
+    velocity_model: Callable,
+    xt: Tensor,
+    t0: Tensor,
+    dt: Tensor,
+    manifold: Manifold,
+    projx: bool = True,
+    proju: bool = True,
+) -> Tensor:
+    r"""Perform an RK4 step on a manifold.
+
+    Args:
+        velocity_model (Callable): the velocity model
+        xt (Tensor): tensor containing the state at time t0
+        t0 (Tensor): the time at which this step is taken
+        dt (Tensor): the step size
+        manifold (Manifold): a manifold object
+        projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+        proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+    Returns:
+        Tensor: tensor containing the state after the step
+    """
+    velocity_fn = lambda x, t: (
+        manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+    )
+    projx_fn = lambda x: manifold.projx(x) if projx else x
+
+    k1 = velocity_fn(xt, t0)
+    k2 = velocity_fn(projx_fn(xt + dt * k1 / 3), t0 + dt / 3)
+    k3 = velocity_fn(projx_fn(xt + dt * (k2 - k1 / 3)), t0 + dt * 2 / 3)
+    k4 = velocity_fn(projx_fn(xt + dt * (k1 - k2 + k3)), t0 + dt)
+
+    return projx_fn(xt + (k1 + 3 * (k2 + k3) + k4) * dt * 0.125)

flow_matching/solver/solver.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+
+from torch import nn, Tensor
+
+
+class Solver(ABC, nn.Module):
+    """Abstract base class for solvers."""
+
+    @abstractmethod
+    def sample(self, x_0: Tensor = None) -> Tensor:
+        ...

flow_matching/solver/utils.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+
+
+def get_nearest_times(time_grid: Tensor, t_discretization: Tensor) -> Tensor:
+    distances = torch.cdist(
+        time_grid.unsqueeze(1),
+        t_discretization.unsqueeze(1),
+        compute_mode="donot_use_mm_for_euclid_dist",
+    )
+    nearest_indices = distances.argmin(dim=1)
+
+    return t_discretization[nearest_indices]

flow_matching/utils/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .categorical_sampler import categorical
+from .model_wrapper import ModelWrapper
+from .utils import expand_tensor_like, gradient, unsqueeze_to_match
+
+__all__ = [
+    "unsqueeze_to_match",
+    "expand_tensor_like",
+    "gradient",
+    "categorical",
+    "ModelWrapper",
+]

flow_matching/utils/categorical_sampler.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+
+
+def categorical(probs: Tensor) -> Tensor:
+    r"""Categorical sampler according to weights in the last dimension of ``probs`` using :func:`torch.multinomial`.
+
+    Args:
+        probs (Tensor): probabilities.
+
+    Returns:
+        Tensor: Samples.
+    """
+
+    return torch.multinomial(probs.flatten(0, -2), 1, replacement=True).view(
+        *probs.shape[:-1]
+    )

flow_matching/utils/manifolds/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .manifold import Euclidean, Manifold
+from .sphere import Sphere
+from .torus import FlatTorus
+from .utils import geodesic
+
+__all__ = [
+    "Euclidean",
+    "Manifold",
+    "Sphere",
+    "FlatTorus",
+    "geodesic",
+]

flow_matching/utils/manifolds/manifold.py

@@ -0,0 +1,93 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import abc
+
+import torch.nn as nn
+from torch import Tensor
+
+
+class Manifold(nn.Module, metaclass=abc.ABCMeta):
+    """A manifold class that contains projection operations and logarithm and exponential maps."""
+
+    @abc.abstractmethod
+    def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+        r"""Computes exponential map :math:`\exp_x(u)`.
+
+        Args:
+            x (Tensor): point on the manifold
+            u (Tensor): tangent vector at point :math:`x`
+
+        Raises:
+            NotImplementedError: if not implemented
+
+        Returns:
+            Tensor: transported point
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+        r"""Computes logarithmic map :math:`\log_x(y)`.
+
+        Args:
+            x (Tensor): point on the manifold
+            y (Tensor): point on the manifold
+
+        Raises:
+            NotImplementedError: if not implemented
+
+        Returns:
+            Tensor: tangent vector at point :math:`x`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def projx(self, x: Tensor) -> Tensor:
+        """Project point :math:`x` on the manifold.
+
+        Args:
+            x (Tensor): point to be projected
+
+        Raises:
+            NotImplementedError: if not implemented
+
+        Returns:
+            Tensor: projected point on the manifold
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def proju(self, x: Tensor, u: Tensor) -> Tensor:
+        """Project vector :math:`u` on a tangent space for :math:`x`.
+
+        Args:
+            x (Tensor): point on the manifold
+            u (Tensor): vector to be projected
+
+        Raises:
+            NotImplementedError: if not implemented
+
+        Returns:
+            Tensor: projected tangent vector
+        """
+        raise NotImplementedError
+
+
+class Euclidean(Manifold):
+    """The Euclidean manifold."""
+
+    def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+        return x + u
+
+    def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+        return y - x
+
+    def projx(self, x: Tensor) -> Tensor:
+        return x
+
+    def proju(self, x: Tensor, u: Tensor) -> Tensor:
+        return u