Add Flow Map Matching for Torch Backend

[lorenzo-consoli]

No description provided.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Copilot]

Pull request overview

This PR introduces Flow Map Matching (FMM) support for the Torch backend by adding a new flow-map neural module and a corresponding training/prediction method, plus tests and example notebooks.

Changes:

• Add MLPFlowMap (flow-map model) and FMM (flow-map matching method) to the Torch backend and register the new method.
• Refactor Torch NN velocity-field internals to use _nn modules dict (and t_encoder key) and add a new TFmFn type.
• Add/adjust test coverage and provide toy notebooks for trying CFM/FMM.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
src/sc_flow/backends/torch/nn/_fm.py	Adds the MLPFlowMap module (flow map conditioned on (s,t) and x).
src/sc_flow/backends/torch/methods/library/_fmm.py	Adds the FMM method implementation (losses + prediction loop).
src/sc_flow/backends/torch/methods/library/__init__.py	Exposes FMM in the library module exports.
src/sc_flow/backends/torch/methods/__init__.py	Registers "fmm" in METHODS_REGISTRY.
src/sc_flow/backends/torch/nn/_vf.py	Renames internal module dict from _vf to _nn and time_encoder key to t_encoder.
src/sc_flow/backends/torch/_types.py	Adds TFmFn callable type for flow-map functions.
src/sc_flow/backends/torch/methods/_base.py	Adds _prepare_latent_state helper to the Torch generative-flow base class.
src/sc_flow/trainer/_trainer.py	Changes logged training step numbering to be 1-based.
src/sc_flow/data/containers/_categorical.py	Plumbs categorical_reps_map through from_pandas.
tests/backends/torch/nn/test_fm.py	Adds unit tests for MLPFlowMap.
tests/backends/torch/methods/library/test_fmm.py	Adds unit tests for FMM (loss dispatch, predict shapes, etc.).
tests/backends/torch/methods/library/test_cfm.py	Updates tests to call _step_fn instead of _compute_loss.
docs/notebooks/try_fmm_toy.ipynb	Adds a toy notebook demonstrating FMM usage.
docs/notebooks/try_cfm_toy.ipynb	Clears execution counts, but still updates outputs.

Comments suppressed due to low confidence (1)

src/sc_flow/backends/torch/methods/init.py:13

• AVAILABLE_METHODS is still Literal["cfm"] even though the registry now includes the new "fmm" method. This will cause type/validation mismatches for consumers that use AVAILABLE_METHODS. Update the literal to include "fmm".

METHODS_REGISTRY = {
    "cfm": CFM,
    "fmm": FMM,
}
AVAILABLE_METHODS = Literal["cfm"]

__all__ = ["TorchBaseMethod", "TorchGenerativeFlow", "CFM", "METHODS_REGISTRY", "AVAILABLE_METHODS"]

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[lorenzo-consoli]

We should make this more generic for consistency models.