Sync from GitHub FlexiBrain main

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/pipeline.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,25 @@
+__pycache__/
+*.py[cod]
+*.so
+*.whl
+*.egg-info/
+build/
+dist/
+.eggs/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+/logs/
+/checkpoints/
+/resolution_runs/
+/outputs/
+/data/
+*.nii
+*.nii.gz
+*.pt
+*.pth
+*.npy
+*.npz
+.DS_Store
+Thumbs.db
+.nfs*

LICENSE

@@ -0,0 +1,12 @@
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+
+Flexibrain is provided for non-commercial research use under the Creative
+Commons Attribution-NonCommercial-ShareAlike 4.0 International License
+(CC BY-NC-SA 4.0).
+
+You may share and adapt the material for non-commercial purposes, provided
+that appropriate attribution is given and adaptations are distributed under
+the same license terms.
+
+Full legal code: https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+Human-readable summary: https://creativecommons.org/licenses/by-nc-sa/4.0/

NOTICE

@@ -0,0 +1,8 @@
+Flexibrain Notice
+
+This repository contains Flexibrain project code plus small source fragments
+needed for the pretrain/downstream path. Third-party references and preserved
+license files are listed in THIRD_PARTY.md and licenses/.
+
+Selected files include reference comments at the top of the source file where
+the implementation is derived from or adapted from prior public projects.

README.md

@@ -1,180 +1,177 @@
----
-license: cc-by-nc-sa-4.0
----
-# Flexibrain
-
-Flexibrain is a voxel-level fMRI representation learning framework for pretraining and downstream classification. It keeps fMRI volumes in a fixed 96 x 96 x 96 input grid, reads each sample's voxel spacing and TR from the NIfTI header, and resizes patch embedding kernels in physical spatial and temporal units before learning with a Mamba-JEPA backbone.
-
-<p align="center">
-  <img src="assets/pipeline.png" width="900" alt="FlexiBrain framework pipeline">
-</p>
-
-## Installation
-
-The code was tested on l40 with Python 3.10, PyTorch 2.1.2, CUDA 12.1, `causal-conv1d`, `mamba-ssm`, and `flash-attn`.
-
-```bash
-conda create -n flexibrain python=3.10
-conda activate flexibrain
-pip install -r requirements.txt
-pip install -e .
-```
-
-Check the CLI:
-
-```bash
-python -m flexibrain --help
-python -m flexibrain pretrain --help
-python -m flexibrain downstream --help
-```
-
-## Data Preparation
-
-Each sample should be a 4D NIfTI file shaped as:
-
-```text
-96 x 96 x 96 x T
-```
-
-Flexibrain uses the NIfTI header to read voxel spacing and TR. If a dataset has missing TR metadata, fix the header before training or pass an explicit fallback with `--default-tr` / `data.default_tr`.
-
-`T_prime` and `tau_seconds` control the selected temporal length:
-
-```text
-kt = round(tau_seconds / TR)
-T_selected = T_prime * kt
-```
-
-The preprocessing script can convert native/T1/MNI-space inputs to 96 x 96 x 96, apply sample-wise global z-score normalization over foreground voxels, and write 4D NIfTI outputs:
-
-```bash
-python data_process.py \
-  --input-root /path/to/input_root \
-  --output-root /path/to/output_root \
-  --spaces all \
-  --groups class0,class1,class2
-```
-
-Expected grouped input layout:
-
-```text
-input_root/
-|-- nativespace/class0/*.nii.gz
-|-- t1space/class0/*.nii.gz
-`-- mnispace/class0/*.nii.gz
-```
-
-If files are not organized by group subfolders, omit `--groups`. For MNI-space inputs, provide `--template-mask` when the default mask is not available.
-
-Pretraining list files contain one NIfTI path per line:
-
-```text
-/path/to/sub-0001_bold.nii.gz
-/path/to/sub-0002_bold.nii.gz
-```
-
-Downstream classification uses the same list format plus a CSV label table:
-
-```csv
-Subject,Group_idx
-003_S_0908,2
-011_S_0002,1
-1001,0
-```
-
-Default label fields are `Subject` and `Group_idx`. `path_id_mode=auto` supports ADNI-style IDs such as `003_S_0908`, ADHD-style filenames, and fallback digit IDs.
-
-## Pretraining
-
-Run from a config:
-
-```bash
-python -m flexibrain pretrain --config configs/pretrain_example.yaml
-```
-
-Or use CLI arguments:
-
-```bash
-python -m flexibrain pretrain \
-  --train-list /path/to/pretrain_train.txt \
-  --val-list /path/to/pretrain_val.txt \
-  --checkpoint-dir ./checkpoints/pretrain/example \
-  --log-dir ./logs/pretrain/example \
-  --embed-dim 512 \
-  --depth 24 \
-  --predictor-depth 2 \
-  --bimamba-type v2 \
-  --if-devide-out \
-  --batch-size 4 \
-  --epochs 30 \
-  --lr 5e-4 \
-  --weight-decay 0.05 \
-  --warmup-epochs 3 \
-  --mask-ratio 0.65 \
-  --grad-accumulation-steps 4 \
-  --t-prime 30 \
-  --tau-seconds 6.0 \
-  --use-amp
-```
-
-Outputs:
-
-```text
-checkpoint_latest.pt
-checkpoint_best.pt
-pretrain_*.log
-```
-
-## Downstream Classification
-
-Run from a config:
-
-```bash
-python -m flexibrain downstream --config configs/downstream_example.yaml
-```
-
-Or use CLI arguments:
-
-```bash
-python -m flexibrain downstream \
-  --train-list /path/to/downstream_train.txt \
-  --val-list /path/to/downstream_val.txt \
-  --test-list /path/to/downstream_test.txt \
-  --csv /path/to/labels.csv \
-  --pretrain-checkpoint /path/to/checkpoint_best.pt \
-  --num-classes 3 \
-  --head-type transformer \
-  --batch-size 8 \
-  --epochs 30 \
-  --lr 1e-5 \
-  --lr-backbone 6e-6 \
-  --lr-head 6e-5 \
-  --checkpoint-dir ./checkpoints/downstream/example \
-  --log-dir ./logs/downstream/example \
-  --use-amp
-```
-
-During downstream training, validation metrics select `downstream_best.pt`. The test set is evaluated once at the end after loading that best validation checkpoint, and the final metrics are written to `test_metrics.json`.
-
-## Configuration
-
-YAML config mirrors the CLI options. Keep private paths in local config files and leave shared configs as portable examples. The provided examples use placeholder paths under `data/`:
-
-```text
-configs/pretrain_example.yaml
-configs/downstream_example.yaml
-```
-
-## Checkpoint Compatibility
-
-The downstream loader can initialize from the original pretraining checkpoint path format:
-
-```text
-/path/to/checkpoint_best.pt
-```
-
-When `use_checkpoint_config: true`, model-shape settings stored in the checkpoint are applied before loading the backbone.
-
-## License
-
-This repository is provided for non-commercial research use under CC BY-NC-SA 4.0. See `LICENSE` and `NOTICE` for license boundaries and preserved notices.
+# Flexibrain
+
+Flexibrain is a voxel-level fMRI representation learning framework for pretraining and downstream classification. It keeps fMRI volumes in a fixed 96 x 96 x 96 input grid, reads each sample's voxel spacing and TR from the NIfTI header, and resizes patch embedding kernels in physical spatial and temporal units before learning with a Mamba-JEPA backbone.
+
+<p align="center">
+  <img src="assets/pipeline.png" width="900" alt="FlexiBrain framework pipeline">
+</p>
+
+## Installation
+
+The code was tested on l40 with Python 3.10, PyTorch 2.1.2, CUDA 12.1, `causal-conv1d`, `mamba-ssm`, and `flash-attn`.
+
+```bash
+conda create -n flexibrain python=3.10
+conda activate flexibrain
+pip install -r requirements.txt
+pip install -e .
+```
+
+Check the CLI:
+
+```bash
+python -m flexibrain --help
+python -m flexibrain pretrain --help
+python -m flexibrain downstream --help
+```
+
+## Data Preparation
+
+Each sample should be a 4D NIfTI file shaped as:
+
+```text
+96 x 96 x 96 x T
+```
+
+Flexibrain uses the NIfTI header to read voxel spacing and TR. If a dataset has missing TR metadata, fix the header before training or pass an explicit fallback with `--default-tr` / `data.default_tr`.
+
+`T_prime` and `tau_seconds` control the selected temporal length:
+
+```text
+kt = round(tau_seconds / TR)
+T_selected = T_prime * kt
+```
+
+The preprocessing script can convert native/T1/MNI-space inputs to 96 x 96 x 96, apply sample-wise global z-score normalization over foreground voxels, and write 4D NIfTI outputs:
+
+```bash
+python data_process.py \
+  --input-root /path/to/input_root \
+  --output-root /path/to/output_root \
+  --spaces all \
+  --groups class0,class1,class2
+```
+
+Expected grouped input layout:
+
+```text
+input_root/
+|-- nativespace/class0/*.nii.gz
+|-- t1space/class0/*.nii.gz
+`-- mnispace/class0/*.nii.gz
+```
+
+If files are not organized by group subfolders, omit `--groups`. For MNI-space inputs, provide `--template-mask` when the default mask is not available.
+
+Pretraining list files contain one NIfTI path per line:
+
+```text
+/path/to/sub-0001_bold.nii.gz
+/path/to/sub-0002_bold.nii.gz
+```
+
+Downstream classification uses the same list format plus a CSV label table:
+
+```csv
+Subject,Group_idx
+003_S_0908,2
+011_S_0002,1
+1001,0
+```
+
+Default label fields are `Subject` and `Group_idx`. `path_id_mode=auto` supports ADNI-style IDs such as `003_S_0908`, ADHD-style filenames, and fallback digit IDs.
+
+## Pretraining
+
+Run from a config:
+
+```bash
+python -m flexibrain pretrain --config configs/pretrain_example.yaml
+```
+
+Or use CLI arguments:
+
+```bash
+python -m flexibrain pretrain \
+  --train-list /path/to/pretrain_train.txt \
+  --val-list /path/to/pretrain_val.txt \
+  --checkpoint-dir ./checkpoints/pretrain/example \
+  --log-dir ./logs/pretrain/example \
+  --embed-dim 512 \
+  --depth 24 \
+  --predictor-depth 2 \
+  --bimamba-type v2 \
+  --if-devide-out \
+  --batch-size 4 \
+  --epochs 30 \
+  --lr 5e-4 \
+  --weight-decay 0.05 \
+  --warmup-epochs 3 \
+  --mask-ratio 0.65 \
+  --grad-accumulation-steps 4 \
+  --t-prime 30 \
+  --tau-seconds 6.0 \
+  --use-amp
+```
+
+Outputs:
+
+```text
+checkpoint_latest.pt
+checkpoint_best.pt
+pretrain_*.log
+```
+
+## Downstream Classification
+
+Run from a config:
+
+```bash
+python -m flexibrain downstream --config configs/downstream_example.yaml
+```
+
+Or use CLI arguments:
+
+```bash
+python -m flexibrain downstream \
+  --train-list /path/to/downstream_train.txt \
+  --val-list /path/to/downstream_val.txt \
+  --test-list /path/to/downstream_test.txt \
+  --csv /path/to/labels.csv \
+  --pretrain-checkpoint /path/to/checkpoint_best.pt \
+  --num-classes 3 \
+  --head-type transformer \
+  --batch-size 8 \
+  --epochs 30 \
+  --lr 1e-5 \
+  --lr-backbone 6e-6 \
+  --lr-head 6e-5 \
+  --checkpoint-dir ./checkpoints/downstream/example \
+  --log-dir ./logs/downstream/example \
+  --use-amp
+```
+
+During downstream training, validation metrics select `downstream_best.pt`. The test set is evaluated once at the end after loading that best validation checkpoint, and the final metrics are written to `test_metrics.json`.
+
+## Configuration
+
+YAML config mirrors the CLI options. Keep private paths in local config files and leave shared configs as portable examples. The provided examples use placeholder paths under `data/`:
+
+```text
+configs/pretrain_example.yaml
+configs/downstream_example.yaml
+```
+
+## Checkpoint Compatibility
+
+The downstream loader can initialize from the original pretraining checkpoint path format:
+
+```text
+/path/to/checkpoint_best.pt
+```
+
+When `use_checkpoint_config: true`, model-shape settings stored in the checkpoint are applied before loading the backbone.
+
+## License
+
+This repository is provided for non-commercial research use under CC BY-NC-SA 4.0. See `LICENSE` and `NOTICE` for license boundaries and preserved notices.

THIRD_PARTY.md

@@ -0,0 +1,10 @@
+# Third-Party References
+
+Flexibrain keeps only the source fragments needed for the Mamba-JEPA pretrain/downstream path.
+
+- Brain-Harmony / BrainHarmonix: https://github.com/hzlab/Brain-Harmony. The Transformer head block in `flexibrain/models/transformer_block.py` is derived from the official Brain-Harmony `libs/flex_transformer.py` design. The official README marks Brain-Harmony as CC BY-NC-SA 4.0. No standalone Brain-Harmony project directory is vendored.
+- 3D Mamba MAE: referenced for the Mamba block factory design. The original project license file is preserved in `licenses/mamba_mae_LICENSE_CC_BY_NC_4.0.txt`.
+- Mamba / mamba_ssm: the minimal Python source needed by the custom Mamba block is included under `mamba_ssm/`; the Apache-2.0 license is preserved in `licenses/mamba2_LICENSE_Apache_2.0.txt`.
+- causal-conv1d: used as a CUDA extension dependency and installed via requirements; BSD-3-Clause license is preserved in `licenses/causal_conv1d_LICENSE_BSD_3_Clause.txt`.
+
+Binary build artifacts (`*.so`, `*.whl`), checkpoints, logs, and datasets are not included in this repository.

configs/downstream_example.yaml

@@ -0,0 +1,45 @@
+pretrain_checkpoint: /path/to/checkpoint_best.pt
+from_scratch: false
+use_checkpoint_config: true
+model:
+  model_type: mamba
+  num_classes: 3
+  head_type: transformer
+  head_depth: 2
+  head_num_heads: 8
+  head_mlp_ratio: 4.0
+  head_proj_drop: 0.1
+  head_drop_path: 0.1
+  mlp_hidden: 512
+  mlp_depth: 4
+  mlp_dropout: 0.1
+  freeze_backbone: false
+data:
+  train_list: data/lists/downstream_train.txt
+  val_list: data/lists/downstream_val.txt
+  test_list: data/lists/downstream_test.txt
+  csv: data/labels.csv
+  id_column: Subject
+  label_column: Group_idx
+  label_mode: multiclass
+  path_id_mode: auto
+  batch_size: 8
+  num_workers: 8
+  T_prime: 30
+  tau_seconds: 6.0
+  default_tr: null
+training:
+  epochs: 30
+  lr: 1.0e-5
+  lr_backbone: 6.0e-6
+  lr_head: 6.0e-5
+  weight_decay: 0.05
+  warmup_epochs: 2
+  grad_accumulation_steps: 2
+  grad_clip: 1.0
+  seed: 42
+  use_amp: true
+logging:
+  log_interval: 20
+  checkpoint_dir: ./checkpoints/downstream/example
+  log_dir: ./logs/downstream/example

configs/pretrain_example.yaml

@@ -0,0 +1,34 @@
+model:
+  model_type: mamba
+  embed_dim: 512
+  depth: 24
+  predictor_depth: 2
+  drop_path_rate: 0.1
+  bimamba_type: v2
+  if_devide_out: true
+  mixer_type: mamba
+  momentum: 0.996
+  final_momentum: 0.9999
+  norm_target: true
+data:
+  train_list: data/lists/pretrain_train.txt
+  val_list: data/lists/pretrain_val.txt
+  batch_size: 4
+  num_workers: 8
+  T_prime: 30
+  tau_seconds: 6.0
+  default_tr: null
+training:
+  epochs: 30
+  lr: 5.0e-4
+  weight_decay: 0.05
+  warmup_epochs: 3
+  mask_ratio: 0.65
+  grad_clip: 1.0
+  grad_accumulation_steps: 4
+  seed: 42
+  use_amp: true
+logging:
+  log_interval: 20
+  checkpoint_dir: ./checkpoints/pretrain/example
+  log_dir: ./logs/pretrain/example

data_process.py

@@ -0,0 +1,535 @@
+#!/usr/bin/env python3
+"""Generic three-space 4D voxel preprocessing with sample-wise global z-score."""
+
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import nibabel as nib
+import numpy as np
+
+try:
+    from nilearn.image import resample_to_img
+except ImportError:
+    resample_to_img = None
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    tqdm = None
+
+
+SPACE_ORDER = ("native", "t1", "mni")
+
+
+@dataclass
+class Counts:
+    seen: int = 0
+    processed: int = 0
+    dry_run: int = 0
+    skipped_existing: int = 0
+    skipped_input_dir: int = 0
+    skipped_non4d: int = 0
+    skipped_no_mask: int = 0
+    skipped_empty_foreground: int = 0
+    failed: int = 0
+
+    def add(self, other: "Counts") -> None:
+        for name in self.__dataclass_fields__:
+            setattr(self, name, getattr(self, name) + getattr(other, name))
+
+
+def parse_target_shape(value: str) -> Tuple[int, int, int]:
+    parts = [part.strip() for part in value.split(",")]
+    if len(parts) != 3:
+        raise argparse.ArgumentTypeError("target shape must be formatted as X,Y,Z")
+    try:
+        shape = tuple(int(part) for part in parts)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("target shape values must be integers") from exc
+    if any(dim <= 0 for dim in shape):
+        raise argparse.ArgumentTypeError("target shape values must be positive")
+    return shape  # type: ignore[return-value]
+
+
+def parse_csv(value: Optional[str]) -> Optional[List[str]]:
+    if value is None:
+        return None
+    items = [item.strip() for item in value.split(",") if item.strip()]
+    return items or None
+
+
+def selected_spaces(value: str) -> List[str]:
+    if value == "all":
+        return list(SPACE_ORDER)
+    return [value]
+
+
+def is_nifti_file(path: Path) -> bool:
+    name = path.name.lower()
+    return path.is_file() and (name.endswith(".nii") or name.endswith(".nii.gz"))
+
+
+def nifti_stem(path: Path) -> str:
+    name = path.name
+    if name.lower().endswith(".nii.gz"):
+        return name[:-7]
+    if name.lower().endswith(".nii"):
+        return name[:-4]
+    return path.stem
+
+
+def clean_output_name(in_file: Path) -> str:
+    return f"{nifti_stem(in_file)}_global_zscore.nii.gz"
+
+
+def iter_nifti_files(input_dir: Path) -> List[Path]:
+    if not input_dir.is_dir():
+        return []
+    return sorted((path for path in input_dir.iterdir() if is_nifti_file(path)), key=lambda p: p.name)
+
+
+def progress(items: Sequence[Path], desc: str) -> Iterable[Path]:
+    if tqdm is None:
+        return items
+    return tqdm(items, desc=desc)
+
+
+def pad_crop_keep_world(
+    img: nib.Nifti1Image,
+    target_xyz: Tuple[int, int, int] = (96, 96, 96),
+    fill_value: float = 0.0,
+) -> Tuple[nib.Nifti1Image, Dict[str, Any]]:
+    """Center pad/crop the first three axes and update affine translation."""
+    data = np.asarray(img.dataobj)
+    affine = img.affine.copy()
+    linear = affine[:3, :3].copy()
+    old_translation = affine[:3, 3].copy()
+    original_ndim = data.ndim
+
+    if data.ndim == 3:
+        x, y, z = data.shape
+        t = 1
+        data = data[..., None]
+    elif data.ndim == 4:
+        x, y, z, t = data.shape
+    else:
+        raise ValueError(f"Expect 3D/4D NIfTI, got shape {data.shape}")
+
+    def split_plan(old: int, new: int) -> Tuple[int, int, str]:
+        if old == new:
+            return 0, 0, "same"
+        if old < new:
+            total = new - old
+            left = total // 2
+            return left, total - left, "pad"
+        total = old - new
+        left = total // 2
+        return left, total - left, "crop"
+
+    px_l, px_r, mode_x = split_plan(x, target_xyz[0])
+    py_l, py_r, mode_y = split_plan(y, target_xyz[1])
+    pz_l, pz_r, mode_z = split_plan(z, target_xyz[2])
+
+    out_data = data
+    if mode_z == "crop":
+        out_data = out_data[:, :, pz_l : z - pz_r, :]
+    if mode_y == "crop":
+        out_data = out_data[:, py_l : y - py_r, :, :]
+    if mode_x == "crop":
+        out_data = out_data[px_l : x - px_r, :, :, :]
+
+    pad_width = (
+        (px_l if mode_x == "pad" else 0, px_r if mode_x == "pad" else 0),
+        (py_l if mode_y == "pad" else 0, py_r if mode_y == "pad" else 0),
+        (pz_l if mode_z == "pad" else 0, pz_r if mode_z == "pad" else 0),
+        (0, 0),
+    )
+    if any(width != (0, 0) for width in pad_width):
+        out_data = np.pad(out_data, pad_width=pad_width, mode="constant", constant_values=fill_value)
+
+    pad_left = np.array(
+        [
+            px_l if mode_x == "pad" else 0,
+            py_l if mode_y == "pad" else 0,
+            pz_l if mode_z == "pad" else 0,
+        ],
+        dtype=float,
+    )
+    crop_left = np.array(
+        [
+            px_l if mode_x == "crop" else 0,
+            py_l if mode_y == "crop" else 0,
+            pz_l if mode_z == "crop" else 0,
+        ],
+        dtype=float,
+    )
+    new_translation = old_translation + linear @ crop_left - linear @ pad_left
+    affine[:3, 3] = new_translation
+
+    if original_ndim == 3:
+        out_data = out_data[..., 0]
+
+    header = img.header.copy()
+    qcode = int(header.get("qform_code", 1)) or 1
+    scode = int(header.get("sform_code", 1)) or 1
+    out_img = nib.Nifti1Image(out_data, affine, header=header)
+    out_img.set_qform(affine, code=qcode)
+    out_img.set_sform(affine, code=scode)
+
+    info = {
+        "old_shape": (x, y, z, t),
+        "new_shape": out_img.shape if len(out_img.shape) == 4 else out_img.shape + (1,),
+        "plan": {
+            "x": (px_l, px_r, mode_x),
+            "y": (py_l, py_r, mode_y),
+            "z": (pz_l, pz_r, mode_z),
+        },
+        "affine_old_t": tuple(old_translation.tolist()),
+        "affine_new_t": tuple(new_translation.tolist()),
+    }
+    return out_img, info
+
+
+def extract_match_key(path: Path) -> str:
+    stem = nifti_stem(path)
+    bids_match = re.search(r"(sub-[A-Za-z0-9]+)", stem, flags=re.IGNORECASE)
+    if bids_match:
+        return bids_match.group(1).lower()
+
+    for token in ("_space-", "_desc-", "_task-", "_run-", "_bold", "_fmri"):
+        idx = stem.lower().find(token)
+        if idx > 0:
+            return stem[:idx].lower()
+    return stem.lower()
+
+
+def best_matching_file(source_file: Path, target_dir: Path) -> Optional[Path]:
+    if not target_dir.is_dir():
+        return None
+
+    key = extract_match_key(source_file)
+    candidates = [path for path in target_dir.rglob("*") if is_nifti_file(path)]
+    if not candidates:
+        return None
+
+    source_stem = nifti_stem(source_file).lower()
+
+    def score(path: Path) -> Tuple[int, int, int, str]:
+        stem = nifti_stem(path).lower()
+        exact = 3 if stem == source_stem else 0
+        key_match = 2 if key and key in stem else 0
+        short_path = -len(str(path))
+        return (exact, key_match, short_path, str(path))
+
+    best = max(candidates, key=score)
+    best_score = score(best)
+    if best_score[0] == 0 and best_score[1] == 0:
+        return None
+    return best
+
+
+def make_output_image(z_data: np.ndarray, template_img: nib.Nifti1Image) -> nib.Nifti1Image:
+    header = template_img.header.copy()
+    qcode = int(header.get("qform_code", 1)) or 1
+    scode = int(header.get("sform_code", 1)) or 1
+    out_img = nib.Nifti1Image(z_data.astype(np.float32), template_img.affine, header=header)
+    out_img.set_data_dtype(np.float32)
+    out_img.set_qform(template_img.affine, code=qcode)
+    out_img.set_sform(template_img.affine, code=scode)
+    return out_img
+
+
+def global_zscore(arr: np.ndarray, background_mask: np.ndarray) -> Optional[np.ndarray]:
+    foreground_mask = ~background_mask
+    if foreground_mask.shape != arr.shape[:3]:
+        raise ValueError(f"mask shape {foreground_mask.shape} does not match image shape {arr.shape[:3]}")
+    if not np.any(foreground_mask):
+        return None
+
+    vals = arr[foreground_mask, :]
+    if vals.size == 0:
+        return None
+
+    mu = float(vals.mean())
+    std = float(vals.std())
+    with np.errstate(invalid="ignore", divide="ignore"):
+        z = (arr - mu) / (std + 1e-6)
+    z[background_mask, :] = 0.0
+    return z.astype(np.float32, copy=False)
+
+
+def load_4d_padded(path: Path, target_shape: Tuple[int, int, int]) -> Tuple[Optional[nib.Nifti1Image], Optional[Dict[str, Any]]]:
+    img = nib.load(str(path))
+    if len(img.shape) != 4:
+        return None, None
+    return pad_crop_keep_world(img, target_xyz=target_shape, fill_value=0.0)
+
+
+def resolve_mask_path(mask_arg: Optional[str], script_dir: Path) -> Optional[Path]:
+    if not mask_arg:
+        return None
+    mask_path = Path(mask_arg)
+    if mask_path.exists():
+        return mask_path
+    script_relative = script_dir / mask_arg
+    if script_relative.exists():
+        return script_relative
+    return mask_path
+
+
+def resample_template_background(mask_img: nib.Nifti1Image, target_img: nib.Nifti1Image) -> np.ndarray:
+    if resample_to_img is None:
+        raise RuntimeError("nilearn is required for template mask resampling but is not installed")
+    try:
+        mask_res = resample_to_img(
+            mask_img,
+            target_img,
+            interpolation="nearest",
+            force_resample=True,
+            copy_header=True,
+        )
+    except TypeError:
+        mask_res = resample_to_img(mask_img, target_img, interpolation="nearest")
+    mask_arr = np.asarray(mask_res.dataobj)
+    if mask_arr.ndim == 4:
+        mask_arr = mask_arr[..., 0]
+    return mask_arr == 0
+
+
+def space_subdir(args: argparse.Namespace, space: str) -> str:
+    return {
+        "native": args.native_subdir,
+        "t1": args.t1_subdir,
+        "mni": args.mni_subdir,
+    }[space]
+
+
+def input_dir_for(args: argparse.Namespace, space: str, group: Optional[str]) -> Path:
+    base = args.input_root / space_subdir(args, space)
+    return base / group if group else base
+
+
+def output_dir_for(args: argparse.Namespace, space: str, group: Optional[str]) -> Path:
+    base = args.output_root / space_subdir(args, space)
+    return base / group if group else base
+
+
+def output_path_for(args: argparse.Namespace, space: str, group: Optional[str], in_file: Path) -> Path:
+    return output_dir_for(args, space, group) / clean_output_name(in_file)
+
+
+def discover_groups(args: argparse.Namespace, space: str) -> List[Optional[str]]:
+    if args.groups is not None:
+        return args.groups
+
+    base = args.input_root / space_subdir(args, space)
+    if not base.is_dir():
+        return [None]
+
+    groups: List[Optional[str]] = []
+    if iter_nifti_files(base):
+        groups.append(None)
+    groups.extend(sorted(path.name for path in base.iterdir() if path.is_dir() and iter_nifti_files(path)))
+    return groups or [None]
+
+
+def process_one_file(
+    in_file: Path,
+    out_file: Path,
+    space: str,
+    group: Optional[str],
+    args: argparse.Namespace,
+    template_mask_img: Optional[nib.Nifti1Image] = None,
+) -> str:
+    new_img, info = load_4d_padded(in_file, args.target_shape)
+    if new_img is None:
+        return "skipped_non4d"
+
+    arr = np.asarray(new_img.dataobj, dtype=np.float32)
+    if arr.ndim != 4:
+        return "skipped_non4d"
+
+    if space == "native":
+        background_mask = (arr == 0).all(axis=3)
+    elif space == "t1":
+        native_dir = input_dir_for(args, "native", group)
+        native_match = best_matching_file(in_file, native_dir)
+        if native_match is None:
+            print(f"[WARN] Native-space mask source not found for {in_file.name} in {native_dir}")
+            return "skipped_no_mask"
+        native_img, _ = load_4d_padded(native_match, args.target_shape)
+        if native_img is None:
+            print(f"[WARN] Native-space mask source is not 4D, skipping: {native_match}")
+            return "skipped_no_mask"
+        native_arr = np.asarray(native_img.dataobj, dtype=np.float32)
+        background_mask = (native_arr == 0).all(axis=3)
+    elif space == "mni":
+        if template_mask_img is None:
+            print(f"[WARN] Template mask unavailable, skipping: {in_file}")
+            return "skipped_no_mask"
+        background_mask = resample_template_background(template_mask_img, new_img)
+    else:
+        raise ValueError(f"Unsupported space: {space}")
+
+    z_data = global_zscore(arr, background_mask)
+    if z_data is None:
+        print(f"[WARN] Empty foreground after masking, skipping: {in_file}")
+        return "skipped_empty_foreground"
+
+    group_label = group if group is not None else "ungrouped"
+    print(
+        f"[AFT] {space}/{group_label} {in_file.name}: "
+        f"{info['old_shape']} -> {new_img.shape}, plan={info['plan']}"
+    )
+    if args.dry_run:
+        print(f"[DRY-RUN] Would save {out_file}")
+        return "dry_run"
+
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+    out_img = make_output_image(z_data, new_img)
+    nib.save(out_img, str(out_file))
+    print(f"[SAVE] {out_file}")
+    return "processed"
+
+
+def increment(counts: Counts, status: str) -> None:
+    if not hasattr(counts, status):
+        raise ValueError(f"Unknown status: {status}")
+    setattr(counts, status, getattr(counts, status) + 1)
+
+
+def process_space_group(
+    space: str,
+    group: Optional[str],
+    args: argparse.Namespace,
+    template_mask_img: Optional[nib.Nifti1Image] = None,
+) -> Counts:
+    counts = Counts()
+    input_dir = input_dir_for(args, space, group)
+    group_label = group if group is not None else "ungrouped"
+
+    if not input_dir.is_dir():
+        print(f"[WARN] Missing input dir, skipping {space}/{group_label}: {input_dir}")
+        counts.skipped_input_dir += 1
+        return counts
+
+    files = iter_nifti_files(input_dir)
+    print(f"[INFO] {space}/{group_label}: found {len(files)} NIfTI file(s) in {input_dir}")
+    for in_file in progress(files, desc=f"{space}/{group_label}"):
+        counts.seen += 1
+        out_file = output_path_for(args, space, group, in_file)
+        if out_file.exists() and not args.overwrite:
+            print(f"[SKIP] Exists: {out_file}")
+            counts.skipped_existing += 1
+            continue
+
+        try:
+            status = process_one_file(in_file, out_file, space, group, args, template_mask_img)
+            increment(counts, status)
+        except Exception as exc:
+            counts.failed += 1
+            print(f"[ERROR] Failed {in_file}: {exc}", file=sys.stderr)
+
+    return counts
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Process native/T1/MNI 4D voxel NIfTIs with sample-wise global z-score."
+    )
+    parser.add_argument(
+        "--input-root",
+        type=Path,
+        default=Path("."),
+        help="Root containing the native, T1, and MNI input subdirectories.",
+    )
+    parser.add_argument(
+        "--output-root",
+        type=Path,
+        default=Path("global_zscore_outputs"),
+        help="Output root for processed full 4D NIfTI files.",
+    )
+    parser.add_argument(
+        "--spaces",
+        choices=["native", "t1", "mni", "all"],
+        default="all",
+        help="Space to process.",
+    )
+    parser.add_argument(
+        "--groups",
+        type=parse_csv,
+        default=None,
+        help="Optional comma-separated group/subfolder list. If omitted, groups are auto-discovered.",
+    )
+    parser.add_argument("--native-subdir", default="nativespace", help="Native-space subdirectory under input/output roots.")
+    parser.add_argument("--t1-subdir", default="t1space", help="T1-space subdirectory under input/output roots.")
+    parser.add_argument("--mni-subdir", default="mnispace", help="MNI/template-space subdirectory under input/output roots.")
+    parser.add_argument(
+        "--target-shape",
+        type=parse_target_shape,
+        default=(96, 96, 96),
+        help="Spatial target shape formatted as X,Y,Z.",
+    )
+    parser.add_argument(
+        "--template-mask",
+        default="MNI152_T1_2mm_Brain_Mask.nii.gz",
+        help="Brain mask for MNI/template-space inputs. Relative paths are also checked next to this script.",
+    )
+    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing outputs.")
+    parser.add_argument("--dry-run", action="store_true", help="Print planned work without writing outputs.")
+    return parser
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    warnings.filterwarnings("default")
+    parser = build_arg_parser()
+    args = parser.parse_args(argv)
+
+    input_root = args.input_root.resolve()
+    output_root = args.output_root.resolve()
+    for space in SPACE_ORDER:
+        input_space_dir = (input_root / space_subdir(args, space)).resolve()
+        output_space_dir = (output_root / space_subdir(args, space)).resolve()
+        if output_space_dir == input_space_dir or input_space_dir in output_space_dir.parents:
+            print(
+                "[ERROR] Output space directory must not be the same as or inside "
+                f"the input space directory: {output_space_dir}",
+                file=sys.stderr,
+            )
+            return 2
+
+    spaces = selected_spaces(args.spaces)
+    template_mask_img = None
+    if "mni" in spaces:
+        mask_path = resolve_mask_path(args.template_mask, Path(__file__).resolve().parent)
+        if mask_path is None or not mask_path.exists():
+            print(f"[WARN] Template mask not found; MNI/template files will be skipped: {args.template_mask}", file=sys.stderr)
+        else:
+            template_mask_img = nib.load(str(mask_path))
+            print(f"[INFO] Loaded template mask: {mask_path}")
+
+    total = Counts()
+    for space in spaces:
+        for group in discover_groups(args, space):
+            counts = process_space_group(space, group, args, template_mask_img=template_mask_img)
+            total.add(counts)
+
+    print(
+        "[SUMMARY] "
+        f"seen={total.seen}, processed={total.processed}, dry_run={total.dry_run}, "
+        f"skipped_existing={total.skipped_existing}, skipped_input_dir={total.skipped_input_dir}, "
+        f"skipped_non4d={total.skipped_non4d}, skipped_no_mask={total.skipped_no_mask}, "
+        f"skipped_empty_foreground={total.skipped_empty_foreground}, failed={total.failed}"
+    )
+    return 1 if total.failed else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

flexibrain/__main__.py

@@ -0,0 +1,4 @@
+from flexibrain.cli import main
+
+if __name__ == "__main__":
+    main()

flexibrain/cli.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict
+
+import torch
+
+from flexibrain.config import load_config
+from flexibrain.engine import DownstreamTrainer, Pretrainer
+from flexibrain.models import build_downstream_model, build_pretrain_model
+
+
+def _add_common(parser):
+    parser.add_argument("--config", default=None)
+    parser.add_argument("--train-list", default=None)
+    parser.add_argument("--val-list", default=None)
+    parser.add_argument("--test-list", default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--num-workers", type=int, default=None)
+    parser.add_argument("--t-prime", type=int, default=None)
+    parser.add_argument("--tau-seconds", type=float, default=None)
+    parser.add_argument("--default-tr", type=float, default=None, help="Fallback TR in seconds when a NIfTI header has no valid TR.")
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--lr", type=float, default=None)
+    parser.add_argument("--weight-decay", type=float, default=None)
+    parser.add_argument("--warmup-epochs", type=int, default=None)
+    parser.add_argument("--grad-accumulation-steps", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--log-interval", type=int, default=None)
+    parser.add_argument("--checkpoint-dir", default=None)
+    parser.add_argument("--log-dir", default=None)
+    parser.add_argument("--local-rank", type=int, default=None)
+    parser.add_argument("--world-size", type=int, default=None)
+    parser.add_argument("--use-amp", action="store_true", default=None)
+    parser.add_argument("--no-use-amp", dest="use_amp", action="store_false")
+    parser.add_argument("--dry-run", action="store_true")
+    parser.set_defaults(use_amp=None)
+
+
+def _add_model(parser):
+    parser.add_argument("--embed-dim", type=int, default=None)
+    parser.add_argument("--depth", type=int, default=None)
+    parser.add_argument("--predictor-depth", type=int, default=None)
+    parser.add_argument("--drop-path-rate", type=float, default=None)
+    parser.add_argument("--bimamba-type", default=None)
+    parser.add_argument("--if-bimamba", action="store_true", default=None)
+    parser.add_argument("--if-devide-out", action="store_true", default=None)
+    parser.add_argument("--no-if-devide-out", dest="if_devide_out", action="store_false")
+    parser.add_argument("--mixer-type", default=None)
+    parser.add_argument("--momentum", type=float, default=None)
+    parser.add_argument("--final-momentum", type=float, default=None)
+
+
+def apply_common(cfg, args):
+    for key in ["train_list", "val_list", "test_list", "batch_size", "num_workers", "epochs", "lr", "weight_decay", "warmup_epochs", "grad_accumulation_steps", "seed", "local_rank", "world_size"]:
+        value = getattr(args, key, None)
+        if value is None:
+            continue
+        target = cfg.data if key in {"train_list", "val_list", "test_list", "batch_size", "num_workers"} else cfg.training
+        setattr(target, key, value)
+    if args.t_prime is not None:
+        cfg.data.T_prime = args.t_prime
+    if args.tau_seconds is not None:
+        cfg.data.tau_seconds = args.tau_seconds
+    if args.default_tr is not None:
+        cfg.data.default_tr = args.default_tr
+    if args.use_amp is not None:
+        cfg.training.use_amp = args.use_amp
+    if args.log_interval is not None:
+        cfg.logging.log_interval = args.log_interval
+    if args.checkpoint_dir is not None:
+        cfg.logging.checkpoint_dir = args.checkpoint_dir
+    if args.log_dir is not None:
+        cfg.logging.log_dir = args.log_dir
+
+
+def apply_model(cfg, args):
+    for key in ["embed_dim", "depth", "predictor_depth", "drop_path_rate", "bimamba_type", "if_bimamba", "if_devide_out", "mixer_type", "momentum", "final_momentum"]:
+        value = getattr(args, key, None)
+        if value is not None:
+            setattr(cfg.model, key, value)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(prog="flexibrain")
+    sub = parser.add_subparsers(dest="command", required=True)
+    pretrain = sub.add_parser("pretrain")
+    _add_common(pretrain)
+    _add_model(pretrain)
+    pretrain.add_argument("--mask-ratio", type=float, default=None)
+    pretrain.add_argument("--grad-clip", type=float, default=None)
+
+    downstream = sub.add_parser("downstream")
+    _add_common(downstream)
+    _add_model(downstream)
+    downstream.add_argument("--csv", default=None)
+    downstream.add_argument("--id-column", default=None)
+    downstream.add_argument("--label-column", default=None)
+    downstream.add_argument("--label-mode", default=None)
+    downstream.add_argument("--path-id-mode", default=None)
+    downstream.add_argument("--pretrain-checkpoint", default=None)
+    downstream.add_argument("--from-scratch", action="store_true")
+    downstream.add_argument("--ignore-checkpoint-config", action="store_true")
+    downstream.add_argument("--num-classes", type=int, default=None)
+    downstream.add_argument("--head-type", choices=["transformer", "avgpool"], default=None)
+    downstream.add_argument("--freeze-backbone", action="store_true")
+    downstream.add_argument("--lr-backbone", type=float, default=None)
+    downstream.add_argument("--lr-head", type=float, default=None)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = load_config(args.config)
+    apply_common(cfg, args)
+    apply_model(cfg, args)
+    if args.command == "pretrain":
+        if args.mask_ratio is not None:
+            cfg.training.mask_ratio = args.mask_ratio
+        if args.grad_clip is not None:
+            cfg.training.grad_clip = args.grad_clip
+        if args.dry_run:
+            model = build_pretrain_model(cfg.model, torch.device("cpu"))
+            print(json.dumps({"config": asdict(cfg), "parameters": sum(p.numel() for p in model.parameters())}, indent=2))
+            return
+        Pretrainer(cfg).fit()
+    elif args.command == "downstream":
+        for key in ["csv", "id_column", "label_column", "label_mode", "path_id_mode"]:
+            value = getattr(args, key, None)
+            if value is not None:
+                setattr(cfg.data, key, value)
+        for key in ["num_classes", "head_type", "freeze_backbone"]:
+            value = getattr(args, key, None)
+            if value is not None:
+                setattr(cfg.model, key, value)
+        if args.lr_backbone is not None:
+            cfg.training.lr_backbone = args.lr_backbone
+        if args.lr_head is not None:
+            cfg.training.lr_head = args.lr_head
+        if args.pretrain_checkpoint is not None:
+            cfg.pretrain_checkpoint = args.pretrain_checkpoint
+        if args.from_scratch:
+            cfg.from_scratch = True
+        if args.ignore_checkpoint_config:
+            cfg.use_checkpoint_config = False
+        if args.dry_run:
+            model = build_downstream_model(cfg.model, torch.device("cpu"), checkpoint_path=cfg.pretrain_checkpoint, from_scratch=cfg.from_scratch, use_checkpoint_config=cfg.use_checkpoint_config)
+            print(json.dumps({"config": asdict(cfg), "parameters": sum(p.numel() for p in model.parameters())}, indent=2))
+            return
+        DownstreamTrainer(cfg).fit()
+
+
+if __name__ == "__main__":
+    main()

flexibrain/config.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import yaml
+
+
+@dataclass
+class ModelConfig:
+    model_type: str = "mamba"
+    embed_dim: int = 512
+    depth: int = 24
+    predictor_depth: int = 2
+    drop_path_rate: float = 0.1
+    rms_norm: bool = False
+    fused_add_norm: bool = True
+    residual_in_fp32: bool = True
+    bimamba_type: str = "v2"
+    if_bimamba: bool = False
+    mixer_type: str = "mamba"
+    if_devide_out: bool = True
+    momentum: float = 0.996
+    final_momentum: float = 0.9999
+    norm_target: bool = True
+    num_heads: int = 8
+    mlp_ratio: float = 4.0
+    head_type: str = "transformer"
+    num_classes: int = 3
+    head_depth: int = 2
+    head_num_heads: int = 8
+    head_mlp_ratio: float = 4.0
+    head_proj_drop: float = 0.1
+    head_drop_path: float = 0.1
+    mlp_hidden: int = 512
+    mlp_depth: int = 4
+    mlp_dropout: float = 0.1
+    freeze_backbone: bool = False
+
+
+@dataclass
+class DataConfig:
+    train_list: str = ""
+    val_list: str = ""
+    test_list: Optional[str] = None
+    csv: Optional[str] = None
+    id_column: str = "Subject"
+    label_column: str = "Group_idx"
+    label_mode: str = "multiclass"
+    path_id_mode: str = "auto"
+    normal_label: int = 2
+    batch_size: int = 8
+    num_workers: int = 8
+    memory_map: bool = True
+    T_prime: int = 30
+    tau_seconds: float = 6.0
+    default_tr: Optional[float] = None
+
+
+@dataclass
+class TrainingConfig:
+    epochs: int = 30
+    lr: float = 5e-4
+    lr_backbone: Optional[float] = None
+    lr_head: Optional[float] = None
+    weight_decay: float = 0.05
+    warmup_epochs: int = 2
+    mask_ratio: float = 0.65
+    grad_clip: float = 1.0
+    grad_accumulation_steps: int = 1
+    seed: int = 42
+    use_amp: bool = False
+    local_rank: int = 0
+    world_size: int = 1
+
+
+@dataclass
+class LoggingConfig:
+    log_interval: int = 20
+    checkpoint_dir: str = "./checkpoints"
+    log_dir: str = "./logs"
+    resume: Optional[str] = None
+
+
+@dataclass
+class RunConfig:
+    model: ModelConfig = field(default_factory=ModelConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    logging: LoggingConfig = field(default_factory=LoggingConfig)
+    pretrain_checkpoint: Optional[str] = None
+    from_scratch: bool = False
+    use_checkpoint_config: bool = True
+
+
+def _update_dataclass(obj, values: dict):
+    for key, value in values.items():
+        if hasattr(obj, key):
+            setattr(obj, key, value)
+
+
+def load_config(path: Optional[str]) -> RunConfig:
+    cfg = RunConfig()
+    if not path:
+        return cfg
+    data = yaml.safe_load(Path(path).read_text()) or {}
+    if "model" in data:
+        _update_dataclass(cfg.model, data["model"] or {})
+    if "data" in data:
+        _update_dataclass(cfg.data, data["data"] or {})
+    if "training" in data:
+        _update_dataclass(cfg.training, data["training"] or {})
+    if "logging" in data:
+        _update_dataclass(cfg.logging, data["logging"] or {})
+    for key in ["pretrain_checkpoint", "from_scratch", "use_checkpoint_config"]:
+        if key in data:
+            setattr(cfg, key, data[key])
+    return cfg
+
+
+def apply_checkpoint_config(model_cfg: ModelConfig, checkpoint_config: dict) -> None:
+    keys = [
+        "model_type", "embed_dim", "depth", "predictor_depth", "drop_path_rate",
+        "rms_norm", "fused_add_norm", "residual_in_fp32", "bimamba_type",
+        "if_bimamba", "mixer_type", "if_devide_out", "momentum", "norm_target",
+        "num_heads", "mlp_ratio",
+    ]
+    for key in keys:
+        if key in checkpoint_config:
+            setattr(model_cfg, key, checkpoint_config[key])

flexibrain/data/__init__.py

@@ -0,0 +1,5 @@
+from flexibrain.data.builders import build_downstream_dataloaders, build_pretrain_dataloaders
+from flexibrain.data.nifti import NiftiTxtDataset
+from flexibrain.data.classification import ClassificationDataset
+
+__all__ = [build_downstream_dataloaders, build_pretrain_dataloaders, NiftiTxtDataset, ClassificationDataset]

flexibrain/data/builders.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+from torch.utils.data import DataLoader, DistributedSampler
+
+from flexibrain.config import DataConfig, TrainingConfig
+from flexibrain.data.nifti import NiftiTxtDataset
+from flexibrain.data.classification import ClassificationDataset, custom_collate_fn as downstream_collate
+from flexibrain.data.collate import custom_collate_fn as pretrain_collate
+
+
+def build_pretrain_dataloaders(data: DataConfig, training: TrainingConfig, rank: int = 0, world_size: int = 1) -> Tuple[DataLoader, DataLoader]:
+    train_set = NiftiTxtDataset(data.train_list, return_torch=True, memory_map=data.memory_map, cache_meta=True, T_prime=data.T_prime, tau_seconds=data.tau_seconds, default_tr=data.default_tr)
+    val_set = NiftiTxtDataset(data.val_list, return_torch=True, memory_map=data.memory_map, cache_meta=True, T_prime=data.T_prime, tau_seconds=data.tau_seconds, default_tr=data.default_tr)
+    train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=True, seed=training.seed) if world_size > 1 else None
+    val_sampler = DistributedSampler(val_set, num_replicas=world_size, rank=rank, shuffle=False, seed=training.seed) if world_size > 1 else None
+    train_loader = DataLoader(train_set, batch_size=data.batch_size, sampler=train_sampler, shuffle=train_sampler is None, num_workers=data.num_workers, pin_memory=True, drop_last=True, collate_fn=pretrain_collate)
+    val_loader = DataLoader(val_set, batch_size=data.batch_size, sampler=val_sampler, shuffle=False, num_workers=data.num_workers, pin_memory=True, drop_last=False, collate_fn=pretrain_collate)
+    return train_loader, val_loader
+
+
+def _classification_dataset(txt_file: Optional[str], data: DataConfig):
+    if not txt_file:
+        return None
+    if not data.csv:
+        raise ValueError("data.csv is required for downstream classification")
+    return ClassificationDataset(
+        txt_files=txt_file,
+        csv_path=data.csv,
+        id_column=data.id_column,
+        label_column=data.label_column,
+        label_mode=data.label_mode,
+        path_id_mode=data.path_id_mode,
+        normal_label=data.normal_label,
+        return_torch=True,
+        memory_map=data.memory_map,
+        cache_meta=True,
+        T_prime=data.T_prime,
+        tau_seconds=data.tau_seconds,
+        default_tr=data.default_tr,
+    )
+
+
+def build_downstream_dataloaders(data: DataConfig, training: TrainingConfig, rank: int = 0, world_size: int = 1):
+    train_set = _classification_dataset(data.train_list, data)
+    val_set = _classification_dataset(data.val_list, data)
+    test_set = _classification_dataset(data.test_list, data)
+    train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=True, seed=training.seed) if world_size > 1 else None
+    val_sampler = DistributedSampler(val_set, num_replicas=world_size, rank=rank, shuffle=False, seed=training.seed) if world_size > 1 else None
+    test_sampler = DistributedSampler(test_set, num_replicas=world_size, rank=rank, shuffle=False, seed=training.seed) if world_size > 1 and test_set is not None else None
+    train_loader = DataLoader(train_set, batch_size=data.batch_size, sampler=train_sampler, shuffle=train_sampler is None, num_workers=data.num_workers, pin_memory=True, drop_last=True, collate_fn=downstream_collate)
+    val_loader = DataLoader(val_set, batch_size=data.batch_size, sampler=val_sampler, shuffle=False, num_workers=data.num_workers, pin_memory=True, drop_last=False, collate_fn=downstream_collate)
+    test_loader = None
+    if test_set is not None:
+        test_loader = DataLoader(test_set, batch_size=data.batch_size, sampler=test_sampler, shuffle=False, num_workers=data.num_workers, pin_memory=True, drop_last=False, collate_fn=downstream_collate)
+    return train_loader, val_loader, test_loader

flexibrain/data/classification.py

@@ -0,0 +1,375 @@
+import numpy as np
+from typing import Dict, Any, List, Optional, Tuple, Union
+import torch
+import re
+import ast
+from pathlib import Path
+import pandas as pd
+
+from flexibrain.data.nifti import NiftiTxtDataset, _read_list_files, _load_nifti, _space_time_units_to_mm_s
+
+# class ClassificationDataset(NiftiTxtDataset):
+#     """
+#     Classification dataset that extends NiftiTxtDataset.
+    
+#     Extracts binary labels from filenames:
+#     - Files containing 'control' -> label 0
+#     - Files containing 'patient' -> label 1
+#     """
+    
+#     def __init__(self, *args, **kwargs):
+#         super().__init__(*args, **kwargs)
+#         self.labels = self._extract_labels()
+    
+#     def _extract_labels(self) -> List[int]:
+#         """Extract binary labels from file paths."""
+#         labels = []
+#         for path in self.paths:
+#             path_str = str(path).lower()
+#             if 'cn' in path_str:
+#                 labels.append(0)
+#             elif 'ad' in path_str:
+#                 labels.append(1)
+#             else:
+#                 raise ValueError(
+#                     f"Cannot determine label for {path}. "
+#                     f"Filename must contain 'control' or 'patient'."
+#                 )
+#         return labels
+    
+#     def __getitem__(self, idx: int) -> Dict:
+#         sample = super().__getitem__(idx)
+#         sample['label'] = self.labels[idx]
+#         return sample
+    
+
+
+
+class ClassificationDataset(NiftiTxtDataset):
+
+    _seven_digits = re.compile(r'(\d{4,8})(?!\d)')
+
+    def __init__(
+        self,
+        *args,
+        csv_path: Union[str, Path],
+        id_column: str = 'Subject',
+        label_column: str = 'Group_idx',
+        label_mode: str = 'multiclass',
+        path_id_mode: str = 'auto',
+        normal_label: int = 2,
+        exclude_labels: Optional[List[int]] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.csv_path = Path(csv_path)
+        self.id_column = id_column
+        self.label_column = label_column
+        self.label_mode = str(label_mode)
+        self.path_id_mode = str(path_id_mode)
+        self.normal_label = int(normal_label)
+        self.exclude_labels = set(int(x) for x in (exclude_labels or []))
+
+        self._df = self._load_csv(self.csv_path, self.id_column, self.label_column)
+        self._id_to_label = self._build_id_to_label(self._df, self.id_column, self.label_column)
+
+        self.labels = self._extract_labels()
+        self.valid_indices = [i for i, label in enumerate(self.labels) if label is not None]
+
+    @staticmethod
+    def _normalize_id(x: Any) -> str:
+        s = str(x).strip()
+        if '_' in s:
+            return s.upper()
+        s = s.lstrip('0')
+        return s if s != '' else '0'
+
+    @classmethod
+    def _load_csv(cls, csv_path: Path, id_column: str, label_column: str) -> pd.DataFrame:
+        df = pd.read_csv(csv_path)
+        df = df.copy()
+        df['_norm_id'] = df[id_column].apply(cls._normalize_id)
+        return df
+
+    @staticmethod
+    def _parse_label(val: Any) -> Union[int, str, np.ndarray, None]:
+
+        if pd.isna(val):
+            return None
+
+        if isinstance(val, (list, tuple, np.ndarray)):
+            return np.asarray(val, dtype=int)
+        
+        try:
+            return int(float(val))
+        except Exception:
+            pass
+
+        if isinstance(val, str):
+            s = val.strip()
+            try:
+                lit = ast.literal_eval(s)
+                if isinstance(lit, (list, tuple, np.ndarray)):
+                    return np.asarray(lit, dtype=int)
+                if isinstance(lit, (int, float)):
+                    return int(lit)
+            except Exception:
+                tokens = [t for t in re.split(r'[,\s]+', s) if t]
+                if all(t.isdigit() for t in tokens) and len(tokens) > 1:
+                    return np.asarray([int(t) for t in tokens], dtype=int)
+            return s
+
+        raise ValueError(f"无法解析标签值: {val!r}")
+
+    @classmethod
+    def _build_id_to_label(cls, df: pd.DataFrame, id_column: str, label_column: str):
+        mapping = {}
+        for _, row in df.iterrows():
+            key = row['_norm_id']
+            lbl = cls._parse_label(row[label_column])
+            if lbl is None:
+                continue
+            if key in mapping:
+                a, b = np.asarray(mapping[key]), np.asarray(lbl)
+            else:
+                mapping[key] = lbl
+        return mapping
+
+    def _extract_path_id(self, name: str) -> str:
+        mode = self.path_id_mode.lower()
+        if mode == 'auto':
+            if 'ADNI_' in name or re.search(r'\d{3}_S_\d{4}', name) or re.search(r'sub-\d+', name, flags=re.IGNORECASE):
+                mode = 'adni'
+            elif 'ADHD_' in name:
+                mode = 'adhd'
+            else:
+                mode = 'digits'
+
+        if mode == 'adni':
+            match = re.search(r'(\d{3}_S_\d{4})', name)
+            if match:
+                return match.group(1).upper()
+            match = re.search(r'sub-(\d+)', name, flags=re.IGNORECASE)
+            if match:
+                return self._normalize_id(match.group(1))
+            raise ValueError(f"Cannot extract ADNI subject id from filename: {name}")
+
+        if mode == 'adhd':
+            match = re.search(r'ADHD_[^_]+_(\d+)_', name)
+            if not match:
+                raise ValueError(f"Cannot extract ADHD subject id from filename: {name}")
+            return self._normalize_id(match.group(1))
+
+        matches = self._seven_digits.findall(name)
+        if not matches:
+            raise ValueError(f"Cannot extract subject id from filename: {name}")
+        return self._normalize_id(matches[-1])
+
+    def _extract_labels(self) -> List[Union[int, str, np.ndarray]]:
+
+        labels: List[Union[int, str, np.ndarray]] = []
+        for path in self.paths:
+            name = Path(path).name  
+            norm_id = self._extract_path_id(name)
+
+            label = self._id_to_label.get(norm_id)
+            if label is None:
+                labels.append(None)
+                continue
+            label = self._convert_label(label)
+
+            labels.append(label)
+        return labels
+
+    def _convert_label(self, label: Union[int, str, np.ndarray]) -> Union[int, np.ndarray]:
+        if self.label_mode in ('multiclass', 'raw', None):
+            if not isinstance(label, np.ndarray) and int(label) in self.exclude_labels:
+                return None
+            return label
+        if self.label_mode == 'binary_control_vs_disease':
+            if isinstance(label, np.ndarray):
+                return np.asarray([0 if int(x) == self.normal_label else 1 for x in label], dtype=int)
+            return 0 if int(label) == self.normal_label else 1
+        if self.label_mode == 'binary_pd_vs_prodromal':
+            if isinstance(label, np.ndarray):
+                converted = []
+                for x in label:
+                    xi = int(x)
+                    if xi in self.exclude_labels:
+                        converted.append(-1)
+                    else:
+                        converted.append(1 if xi == 1 else 0)
+                return np.asarray(converted, dtype=int)
+            label_i = int(label)
+            if label_i in self.exclude_labels:
+                return None
+            # PPMI Group_idx: 0=Prodromal, 1=PD, 2=Control.
+            return 1 if label_i == 1 else 0
+        if self.label_mode == 'binary_gender':
+            def to_gender(v: Any) -> int:
+                s = str(v).strip().upper()
+                if s in {'M', 'MALE', '1'}:
+                    return 1
+                if s in {'F', 'FEMALE', '0'}:
+                    return 0
+                raise ValueError(f"Unknown gender label: {v!r}")
+
+            if isinstance(label, np.ndarray):
+                return np.asarray([to_gender(x) for x in label], dtype=int)
+            return to_gender(label)
+        if self.label_mode == 'binary_gender_abide':
+            def to_gender_abide(v: Any) -> int:
+                s = str(v).strip().upper()
+                if s in {'0', 'M', 'MALE'}:
+                    return 1
+                if s in {'1', 'F', 'FEMALE'}:
+                    return 0
+                raise ValueError(f"Unknown ABIDE gender label: {v!r}")
+
+            if isinstance(label, np.ndarray):
+                return np.asarray([to_gender_abide(x) for x in label], dtype=int)
+            return to_gender_abide(label)
+        raise ValueError(f"Unknown label_mode: {self.label_mode}")
+
+    def __len__(self) -> int:
+        return len(self.valid_indices)
+
+    def __getitem__(self, idx: int) -> Dict:
+        raw_idx = self.valid_indices[idx]
+        sample = super().__getitem__(raw_idx)
+        sample['label'] = self.labels[raw_idx]
+        return sample
+
+
+def custom_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Custom collate function to handle:
+    1. nibabel headers and other non-tensor objects
+    2. Variable-length time dimensions (due to different TR values)
+
+    For variable-length data, we pad to the maximum length in the batch.
+    """
+    # Separate tensor/array fields from non-collatable fields
+    tensor_fields = ['data', 'affine']
+    scalar_fields = ['tr', 'subject_idx', 'T_selected', 'T_prime', 'tau_seconds', 'label']
+    tuple_fields = ['voxel']
+    object_fields = ['header', 'path']  # These won't be collated
+
+    collated = {}
+
+    # Handle tensor/array fields with padding for variable-length data
+    for field in tensor_fields:
+        if field in batch[0]:
+            values = [item[field] for item in batch]
+
+            if field == 'data':
+                # Data has variable time dimension due to different TR values
+                # Pad all to the maximum time length
+                max_t = max(v.shape[-1] if len(v.shape) >= 4 else 1 for v in values)
+
+                padded_values = []
+                for v in values:
+                    if len(v.shape) >= 4 and v.shape[-1] < max_t:
+                        # Pad in time dimension (last dimension)
+                        pad_amount = max_t - v.shape[-1]
+                        if isinstance(v, torch.Tensor):
+                            v = torch.nn.functional.pad(v, (0, pad_amount), mode='constant', value=0)
+                        else:
+                            v = np.pad(v, ((0, 0), (0, 0), (0, 0), (0, pad_amount)), mode='constant', constant_values=0)
+                    padded_values.append(v)
+
+                # Convert to tensor and stack
+                if isinstance(padded_values[0], torch.Tensor):
+                    collated[field] = torch.stack(padded_values)
+                else:
+                    collated[field] = torch.from_numpy(np.stack(padded_values))
+            else:
+                # Affine matrices should all be the same size (4x4)
+                if isinstance(values[0], torch.Tensor):
+                    collated[field] = torch.stack(values)
+                else:
+                    collated[field] = torch.from_numpy(np.stack(values))
+
+    # Handle scalar fields
+    for field in scalar_fields:
+        if field in batch[0]:
+            values = [item[field] for item in batch]
+            if isinstance(values[0], (int, float)):
+                collated[field] = torch.tensor(values)
+            else:
+                collated[field] = values
+
+    # Handle tuple fields (like voxel sizes)
+    for field in tuple_fields:
+        if field in batch[0]:
+            collated[field] = [item[field] for item in batch]
+
+    # Handle object fields (keep as lists)
+    for field in object_fields:
+        if field in batch[0]:
+            collated[field] = [item[field] for item in batch]
+
+    return collated
+
+def prepare_batch_data(batch: Dict, device: torch.device) -> Tuple[torch.Tensor, Dict, np.ndarray, torch.Tensor, Optional[torch.Tensor]]:
+    """Prepare batch data for model forward pass.
+
+    Returns:
+        x: Input tensor (B, 96, 96, 96, T_max)
+        meta: Dict {batch_index: {"voxel": (vx, vy, vz), "tr": float}}
+        orig_Ts: Array of original time steps
+        labels: Classification labels
+        affines: Affine matrices or None
+    """
+    # Data is already padded and stacked by custom_collate_fn
+    x = batch['data'].to(device, dtype=torch.float32)
+
+    # Build meta dict - use batch index as key
+    batch_size = x.shape[0]
+    voxels = batch['voxel']
+    trs = batch['tr']
+
+    # Convert trs to numpy if needed
+    if isinstance(trs, torch.Tensor):
+        trs = trs.cpu().numpy()
+    elif isinstance(trs, list):
+        trs = np.array(trs)
+
+    meta = {}
+    for i in range(batch_size):
+        # Get voxel
+        if isinstance(voxels, list):
+            voxel = voxels[i] if isinstance(voxels[i], tuple) else tuple(voxels[i])
+        else:
+            print("voxels is empty")
+            voxel = (2.0, 2.0, 2.0)  # Default voxel size
+
+        # Get TR
+        if isinstance(trs, np.ndarray):
+            tr = float(trs[i])
+        elif isinstance(trs, list):
+            tr = float(trs[i])
+        else:
+            print("trs is empty")
+            tr = 2.0  # Default TR
+
+        meta[i] = {"voxel": voxel, "tr": tr}
+
+    # Get original time steps (T_selected from dataset)
+    orig_Ts = batch.get('T_selected', x.shape[-1])
+    if isinstance(orig_Ts, torch.Tensor):
+        orig_Ts = orig_Ts.cpu().numpy()
+    elif isinstance(orig_Ts, list):
+        orig_Ts = np.array(orig_Ts)
+
+    # Handle labels
+    labels = batch['label']
+    if isinstance(labels, torch.Tensor):
+        labels = labels.to(device, dtype=torch.long)
+    else:
+        labels = torch.tensor(labels, dtype=torch.long, device=device)
+
+    # Get affines if available
+    affines = batch['affine'].to(device, dtype=torch.float32) if 'affine' in batch else None
+
+    return x, meta, orig_Ts, labels, affines

flexibrain/data/collate.py

@@ -0,0 +1,119 @@
+from typing import Dict, List, Any, Tuple, Optional
+import torch
+import numpy as np
+
+
+def custom_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Custom collate function to handle:
+    1. nibabel headers and other non-tensor objects
+    2. Variable-length time dimensions (due to different TR values)
+
+    For variable-length data, we pad to the maximum length in the batch.
+    """
+    # Separate tensor/array fields from non-collatable fields
+    tensor_fields = ['data', 'affine']
+    scalar_fields = ['tr', 'subject_idx', 'T_selected', 'T_prime', 'tau_seconds']
+    tuple_fields = ['voxel']
+    object_fields = ['header', 'path']  # These won't be collated
+
+    collated = {}
+
+    # Handle tensor/array fields with padding for variable-length data
+    for field in tensor_fields:
+        if field in batch[0]:
+            values = [item[field] for item in batch]
+
+            if field == 'data':
+                # Data has variable time dimension due to different TR values
+                # Pad all to the maximum time length
+                max_t = max(v.shape[-1] if len(v.shape) >= 4 else 1 for v in values)
+
+                padded_values = []
+                for v in values:
+                    if len(v.shape) >= 4 and v.shape[-1] < max_t:
+                        # Pad in time dimension (last dimension)
+                        pad_amount = max_t - v.shape[-1]
+                        if isinstance(v, torch.Tensor):
+                            v = torch.nn.functional.pad(v, (0, pad_amount), mode='constant', value=0)
+                        else:
+                            v = np.pad(v, ((0, 0), (0, 0), (0, 0), (0, pad_amount)), mode='constant', value=0)
+                    padded_values.append(v)
+
+                # Convert to tensor and stack
+                if isinstance(padded_values[0], torch.Tensor):
+                    collated[field] = torch.stack(padded_values)
+                else:
+                    collated[field] = torch.from_numpy(np.stack(padded_values))
+            else:
+                # Affine matrices should all be the same size (4x4)
+                if isinstance(values[0], torch.Tensor):
+                    collated[field] = torch.stack(values)
+                else:
+                    collated[field] = torch.from_numpy(np.stack(values))
+
+    # Handle scalar fields
+    for field in scalar_fields:
+        if field in batch[0]:
+            values = [item[field] for item in batch]
+            if isinstance(values[0], (int, float)):
+                collated[field] = torch.tensor(values)
+            else:
+                collated[field] = values
+
+    # Handle tuple fields (like voxel sizes)
+    for field in tuple_fields:
+        if field in batch[0]:
+            collated[field] = [item[field] for item in batch]
+
+    # Handle object fields (keep as lists)
+    for field in object_fields:
+        if field in batch[0]:
+            collated[field] = [item[field] for item in batch]
+
+    return collated
+
+
+def prepare_batch_data(batch: Dict, device: torch.device) -> Tuple[torch.Tensor, Dict, np.ndarray, Optional[torch.Tensor]]:
+    """
+    Prepare batch data for model forward pass.
+
+    Returns:
+        x: Input tensor (B, 96, 96, 96, T_max)
+        meta: Dict {subject_idx: {"voxel": (vx, vy, vz), "tr": float}}
+        orig_Ts: Array of original time steps
+        affines: Affine matrices or None
+    """
+    # Move data to device
+    x = batch['data'].to(device, dtype=torch.float32)
+
+    # Build meta dict: {batch_index: {"voxel": (vx, vy, vz), "tr": float}}
+    subject_idxs = batch['subject_idx'].cpu().numpy()
+    voxels = batch['voxel']  # List of tuples or tensor
+    trs = batch['tr'].cpu().numpy() if isinstance(batch['tr'], torch.Tensor) else batch['tr']
+
+    meta = {}
+    for i, subject_idx in enumerate(subject_idxs):
+        # Handle voxel format
+        if isinstance(voxels, (list, tuple)):
+            voxel = voxels[i]
+        else:
+            voxel = tuple(voxels[i].cpu().numpy()) if isinstance(voxels[i], torch.Tensor) else voxels[i]
+
+        tr = float(trs[i])
+        # Use batch index (i) as key, not subject_idx
+        meta[i] = {"voxel": voxel, "tr": tr}
+
+    # Get original time steps (number of frames, not TR)
+    # T_selected is the actual number of time frames selected by the dataset
+    # Do NOT use 'tr' (time resolution in seconds) as it will cause incorrect T_pad calculation
+    if 'T_selected' in batch:
+        orig_Ts = batch['T_selected'].cpu().numpy() if isinstance(batch['T_selected'], torch.Tensor) else batch['T_selected']
+    else:
+        # Fallback: use actual data time dimension if T_selected is not available
+        orig_Ts = np.array([x.shape[-1] for x in batch['data']])
+
+    # Get affines if available
+    affines = batch['affine'].to(device, dtype=torch.float32) if 'affine' in batch else None
+
+    return x, meta, orig_Ts, affines

flexibrain/data/nifti.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+import os
+import glob
+import random
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import nibabel as nib
+import torch
+from torch.utils.data import Dataset
+
+
+class fMRIDataset(Dataset):
+    def __init__(self, 
+                 data_root, datasets, split_suffixes, crop_length=40, downstream=False):
+
+        self.file_paths = []
+        self.crop_length = crop_length
+        self.downstream = downstream
+
+        for dataset_name in datasets:
+            for suffix in split_suffixes:
+                folder_name = f"{dataset_name}_{suffix}"
+                folder_path = os.path.join(data_root, folder_name)
+                if not os.path.exists(folder_path):
+                    print(f"Warning: Folder not found: {folder_path}")
+                    continue
+
+                for root, dirs, files in os.walk(folder_path):
+                    npz_files = glob.glob(os.path.join(root, "*.npz"))
+                    if len(npz_files) > 1:
+                        sample_size = max(1, int(len(npz_files))) 
+                        npz_files = random.sample(npz_files, sample_size)
+                    self.file_paths.extend(npz_files)
+
+        print(f"Dataset loaded. Total files found: {len(self.file_paths)}")
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        file_path = self.file_paths[idx]
+        try:
+            with np.load(file_path) as data_file:
+                key = list(data_file.keys())[0]
+                fmri_data = data_file[key] 
+                fmri_data = fmri_data.astype(np.float32)
+        except Exception as e:
+            print(f"Error loading file {file_path}: {e}")
+            return None
+
+        total_time_frames = fmri_data.shape[-1]
+        if total_time_frames > self.crop_length:
+            start_idx = np.random.randint(0, total_time_frames - self.crop_length + 1)
+            end_idx = start_idx + self.crop_length
+            cropped_data = fmri_data[..., start_idx:end_idx]
+        else:
+            cropped_data = fmri_data[..., :self.crop_length]
+
+        data_tensor = torch.from_numpy(cropped_data)
+        data_tensor = data_tensor.permute(3, 0, 1, 2)
+
+        return data_tensor
+
+
+def _read_list_files(txt_files: Union[str, Path, Sequence[Union[str, Path]]]) -> List[Path]:
+    """Read one or many .txt files and collect absolute paths listed in them.
+
+    Each line should contain a path to a .nii or .nii.gz file. Empty lines and lines
+    starting with '#' are ignored. Paths are expanded and normalized to absolute Paths.
+    """
+    if isinstance(txt_files, (str, Path)):
+        txt_files = [txt_files]
+    paths: List[Path] = []
+    for f in txt_files:  # type: ignore[assignment]
+        f = Path(f)
+        if not f.exists():
+            raise FileNotFoundError(f"List file not found: {f}")
+        for line in f.read_text().splitlines():
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            p = Path(os.path.expanduser(line)).resolve()
+            # allow relative paths inside list files (relative to the list file dir)
+            if not p.exists():
+                p = (f.parent / line).resolve()
+            if not p.exists():
+                raise FileNotFoundError(f"Path from list file does not exist: {line} (resolved: {p})")
+            if p.suffix not in {".nii", ".gz"} and not str(p).endswith(".nii.gz"):
+                raise ValueError(f"Not a NIfTI file: {p}")
+            paths.append(p)
+    # deduplicate while preserving order
+    seen = set()
+    deduped = []
+    for p in paths:
+        if p not in seen:
+            deduped.append(p)
+            seen.add(p)
+    return deduped
+
+
+def _space_time_units_to_mm_s(header: nib.nifti1.Nifti1Header) -> Tuple[Tuple[float, float, float], float]:
+    """Return (vx, vy, vz) in millimeters and TR in seconds from a NIfTI header.
+
+    Uses header.get_zooms() and header.get_xyzt_units(). Safely handles cases with
+    missing time dimension or unusual units.
+    """
+    zooms = header.get_zooms()
+    # space-time units, e.g. ("mm", "sec")
+    space_u, time_u = header.get_xyzt_units()
+
+    # Spatial voxel sizes
+    vx, vy, vz = (zooms + (1.0, 1.0, 1.0, 1.0))[:3]
+    # Convert to mm if needed
+    if space_u == "m":
+        vx, vy, vz = vx * 1000.0, vy * 1000.0, vz * 1000.0
+    elif space_u in ("mm", None, "unknown"):
+        pass
+    else:
+        # Fallback: assume values already in mm
+        pass
+
+    # Temporal resolution (TR)
+    tr = 0.0
+    if len(zooms) >= 4:
+        tr = float(zooms[3])
+        if time_u == "msec":
+            tr = tr / 1000.0
+        elif time_u in ("usec", "microsec"):
+            tr = tr / 1e6
+        elif time_u in ("sec", None, "unknown"):
+            pass
+        else:
+            # Unknown -> leave as-is
+            pass
+    return (float(vx), float(vy), float(vz)), float(tr)
+
+
+def _load_nifti(path: Union[str, Path], mmap: bool = True) -> Tuple[np.ndarray, np.ndarray, nib.nifti1.Nifti1Header]:
+    try:
+        img = nib.load(str(path), mmap=mmap)
+        data = img.get_fdata(dtype=np.float32)
+        affine = img.affine.copy()
+        header = img.header.copy()
+        return data, affine, header
+    except Exception as e:
+        # Return None to signal invalid file
+        return None, None, None
+
+
+class NiftiTxtDataset(Dataset):
+    """Dataset that loads NIfTI volumes listed in one or more .txt files.
+
+    Each item returns a dict with:
+      - 'data': np.ndarray (from get_fdata())
+      - 'affine': np.ndarray (4x4)
+      - 'header': nibabel header
+      - 'voxel': (vx, vy, vz) in millimeters
+      - 'tr': float, seconds (0.0 if not present)
+      - 'path': pathlib.Path to the NIfTI file
+      - 'subject_idx': integer index inside this dataset
+      - 'T_selected': int, number of time frames selected based on T_prime and tau_seconds
+
+    Parameters
+    ----------
+    txt_files: str | Path | Sequence[str|Path]
+        One or more text files containing absolute (or relative) paths to NIfTI files.
+    transform: Optional[callable]
+        Optional transform applied to the sample dict (after loading).
+    return_torch: bool
+        If True, converts 'data' and 'affine' to torch tensors.
+    memory_map: bool
+        If True, enables nibabel's memory mapping. Disable to force full load into RAM.
+    cache_meta: bool
+        If True, caches voxel/TR in memory to avoid recomputing for repeated access.
+    T_prime: Optional[int]
+        Target number of time patches after TAPE (Time-to-space patch embedding).
+        If provided, dataset will automatically select appropriate time frames to ensure
+        all samples have the same T_prime after TAPE processing.
+        Formula: T_selected = T_prime * tau_seconds / TR
+    tau_seconds: float
+        Time window in seconds for TAPE kernel (default: 6.0).
+        Used to calculate T_selected when T_prime is specified.
+    """
+
+    def __init__(
+        self,
+        txt_files: Union[str, Path, Sequence[Union[str, Path]]],
+        transform: Optional[callable] = None,
+        return_torch: bool = False,
+        memory_map: bool = True,
+        cache_meta: bool = True,
+        T_prime: Optional[int] = None,
+        tau_seconds: float = 6.0,
+        default_tr: Optional[float] = None,
+    ) -> None:
+        super().__init__()
+        self.paths: List[Path] = _read_list_files(txt_files)
+        if len(self.paths) == 0:
+            raise ValueError("No NIfTI paths found in the provided list files.")
+        self.transform = transform
+        self.return_torch = bool(return_torch)
+        self.memory_map = bool(memory_map)
+        self.cache_meta = bool(cache_meta)
+        self.T_prime = T_prime
+        self.tau_seconds = float(tau_seconds)
+        self.default_tr = float(default_tr) if default_tr is not None else None
+        if self.default_tr is not None and self.default_tr <= 0:
+            raise ValueError("default_tr must be positive when provided")
+        self._meta_cache: Dict[int, Tuple[Tuple[float, float, float], float]] = {}
+
+    def __len__(self) -> int:
+        return len(self.paths)
+
+    def _get_meta(self, idx: int, header: Optional[nib.nifti1.Nifti1Header] = None) -> Tuple[Tuple[float, float, float], float]:
+        if self.cache_meta and idx in self._meta_cache:
+            return self._meta_cache[idx]
+        if header is None:
+            _, _, header = _load_nifti(self.paths[idx], mmap=self.memory_map)
+        voxel, tr = _space_time_units_to_mm_s(header)
+        voxel = tuple(float(v) for v in voxel)
+        if any((not np.isfinite(v)) or v <= 0 for v in voxel):
+            raise ValueError(f"Invalid voxel spacing for {self.paths[idx]}: {voxel}")
+        tr = float(tr)
+        if (not np.isfinite(tr)) or tr <= 0:
+            if self.default_tr is None:
+                raise ValueError(
+                    f"Invalid or missing TR for {self.paths[idx]}: {tr}. "
+                    "Set data.default_tr or pass --default-tr to use an explicit fallback."
+                )
+            tr = self.default_tr
+        if self.cache_meta:
+            self._meta_cache[idx] = (voxel, tr)
+        return voxel, tr
+
+    def _calculate_T_selected(self, tr: float, T_total: int) -> int:
+        """
+        Calculate the number of time frames to select based on T_prime and tau_seconds.
+
+        Formula:
+            kt = round(tau_seconds / tr)  # kernel size in time dimension
+            T_selected = T_prime * kt
+
+        This ensures that after TAPE (Time-to-space patch embedding), all samples
+        will have the same number of time patches (T_prime).
+
+        Args:
+            tr: Temporal resolution (TR) in seconds
+            T_total: Total number of time frames available in the data
+
+        Returns:
+            T_selected: Number of time frames to use (min with T_total)
+        """
+        if self.T_prime is None:
+            return T_total
+        if tr <= 0:
+            raise ValueError("TR must be positive when T_prime is set")
+
+        # Calculate kernel size in time dimension
+        kt = max(1, round(self.tau_seconds / tr))
+
+        # Calculate required time frames to get T_prime patches
+        T_selected = self.T_prime * kt
+
+        # Ensure we don't exceed available data
+        T_selected = min(T_selected, T_total)
+
+        return T_selected
+
+    def __getitem__(self, idx: int) -> Dict:
+        # Try to load file, skip to next valid file if current is invalid
+        attempt = 0
+        max_attempts = len(self.paths)
+
+        while attempt < max_attempts:
+            current_idx = (idx + attempt) % len(self.paths)
+            p = self.paths[current_idx]
+            data, affine, header = _load_nifti(p, mmap=self.memory_map)
+
+            # If file is valid, process it
+            if data is not None:
+                voxel, tr = self._get_meta(current_idx, header)
+
+                # # 检测到tr=2或者1.96
+                # if not (np.isclose(tr, 2.0, atol=1e-2) or np.isclose(tr, 1.96, atol=1e-2)):
+                #     attempt += 1
+                #     continue
+                # print(f"TR is {tr} for {p}")
+
+                # Calculate T_selected based on T_prime and tau_seconds
+                T_total = data.shape[3] if len(data.shape) >= 4 else 1
+                T_selected = self._calculate_T_selected(tr, T_total)
+
+                # Slice data to T_selected frames
+                if len(data.shape) >= 4 and T_selected < T_total:
+                    data = data[..., :T_selected]
+                
+                sample = {
+                    "data": torch.from_numpy(data) if self.return_torch else data,
+                    "affine": torch.from_numpy(affine) if self.return_torch else affine,
+                    "header": header,
+                    "voxel": voxel,
+                    "tr": tr,
+                    "path": p,
+                    "subject_idx": current_idx,
+                    "T_selected": T_selected,
+                    "T_prime": self.T_prime,
+                    "tau_seconds": self.tau_seconds,
+                }
+                if self.transform is not None:
+                    sample = self.transform(sample)
+                return sample
+
+            # Try next file if current one is invalid
+            attempt += 1
+
+        # If all files are invalid, raise error
+        raise RuntimeError(f"Could not find any valid file starting from index {idx}")
+
+    def meta_dict(self) -> Dict[int, Dict[str, Union[Tuple[float, float, float], float]]]:
+        """Return {subject_idx: {"voxel": (vx,vy,vz), "tr": tr}} for the whole dataset."""
+        meta: Dict[int, Dict[str, Union[Tuple[float, float, float], float]]] = {}
+        for i, p in enumerate(self.paths):
+            if self.cache_meta and i in self._meta_cache:
+                voxel, tr = self._meta_cache[i]
+            else:
+                # read header cheaply without loading full data
+                img = nib.load(str(p), mmap=True)
+                voxel, tr = _space_time_units_to_mm_s(img.header)
+                if self.cache_meta:
+                    self._meta_cache[i] = (voxel, tr)
+            meta[i] = {"voxel": voxel, "tr": tr}
+        return meta
+
+
+
+
+def build_train_val_from_lists(
+    train_txts: Union[str, Path, Sequence[Union[str, Path]]],
+    val_txts: Union[str, Path, Sequence[Union[str, Path]]],
+    *,
+    return_torch: bool = False,
+    memory_map: bool = True,
+    T_prime: Optional[int] = None,
+    tau_seconds: float = 6.0,
+) -> Tuple[NiftiTxtDataset, NiftiTxtDataset, Dict[str, Dict[int, Dict[str, Union[Tuple[float, float, float], float]]]]]:
+    """Convenience helper to create train/val datasets and collect their meta dicts.
+
+    Parameters
+    ----------
+    train_txts, val_txts: str | Path | Sequence[str|Path]
+        Text files containing paths to NIfTI files
+    return_torch: bool
+        If True, converts data and affine to torch tensors
+    memory_map: bool
+        If True, enables nibabel's memory mapping
+    T_prime: Optional[int]
+        Target number of time patches after TAPE. If provided, dataset will automatically
+        select appropriate time frames to ensure all samples have the same T_prime.
+    tau_seconds: float
+        Time window in seconds for TAPE kernel (default: 6.0)
+
+    Returns
+    -------
+    train_set, val_set, meta_all
+      where meta_all = {"train": {...}, "val": {...}}
+    """
+    train_set = NiftiTxtDataset(
+        train_txts,
+        return_torch=return_torch,
+        memory_map=memory_map,
+        T_prime=T_prime,
+        tau_seconds=tau_seconds,
+    )
+    val_set = NiftiTxtDataset(
+        val_txts,
+        return_torch=return_torch,
+        memory_map=memory_map,
+        T_prime=T_prime,
+        tau_seconds=tau_seconds,
+    )
+    meta_all = {"train": train_set.meta_dict(), "val": val_set.meta_dict()}
+    return train_set, val_set, meta_all
+
+

flexibrain/distributed/__init__.py

@@ -0,0 +1,20 @@
+import os
+import torch.distributed as dist
+
+
+def setup_distributed(rank: int, world_size: int) -> None:
+    """Setup distributed training."""
+    if world_size > 1:
+        os.environ['MASTER_ADDR'] = os.environ.get('MASTER_ADDR', 'localhost')
+        os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '12355')
+        dist.init_process_group(
+            backend='nccl',
+            rank=rank,
+            world_size=world_size
+        )
+
+
+def cleanup_distributed() -> None:
+    """Cleanup distributed training."""
+    if dist.is_available() and dist.is_initialized():
+        dist.destroy_process_group()

flexibrain/engine/__init__.py

@@ -0,0 +1,4 @@
+from flexibrain.engine.pretrainer import Pretrainer
+from flexibrain.engine.downstream_trainer import DownstreamTrainer
+
+__all__ = [Pretrainer, DownstreamTrainer]

flexibrain/engine/downstream_trainer.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+
+from flexibrain.config import RunConfig
+from flexibrain.data import build_downstream_dataloaders
+from flexibrain.data.classification import prepare_batch_data
+from flexibrain.models import build_downstream_model
+from flexibrain.utils.logging import setup_logger
+from flexibrain.utils.seed import set_seed
+
+
+class DownstreamTrainer:
+    def __init__(self, cfg: RunConfig):
+        self.cfg = cfg
+        self.rank = cfg.training.local_rank
+        self.device = torch.device(f"cuda:{self.rank}" if torch.cuda.is_available() else "cpu")
+        self.logger = setup_logger("downstream", cfg.logging.log_dir, rank=self.rank)
+
+    def build(self):
+        set_seed(self.cfg.training.seed)
+        self.model = build_downstream_model(
+            self.cfg.model,
+            self.device,
+            logger=self.logger,
+            checkpoint_path=self.cfg.pretrain_checkpoint,
+            from_scratch=self.cfg.from_scratch,
+            use_checkpoint_config=self.cfg.use_checkpoint_config,
+        )
+        self.train_loader, self.val_loader, self.test_loader = build_downstream_dataloaders(self.cfg.data, self.cfg.training, rank=self.rank, world_size=self.cfg.training.world_size)
+        if self.cfg.training.lr_backbone is not None or self.cfg.training.lr_head is not None:
+            backbone_params = list(self.model.backbone.parameters())
+            head_params = [p for n, p in self.model.named_parameters() if not n.startswith("backbone.")]
+            self.optimizer = optim.AdamW([
+                {"params": backbone_params, "lr": self.cfg.training.lr_backbone or self.cfg.training.lr},
+                {"params": head_params, "lr": self.cfg.training.lr_head or self.cfg.training.lr},
+            ], weight_decay=self.cfg.training.weight_decay)
+        else:
+            self.optimizer = optim.AdamW(self.model.parameters(), lr=self.cfg.training.lr, weight_decay=self.cfg.training.weight_decay)
+        self.use_amp = bool(self.cfg.training.use_amp and self.device.type == "cuda")
+        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)
+        total_steps = max(1, len(self.train_loader) * self.cfg.training.epochs)
+        warmup_steps = max(1, len(self.train_loader) * self.cfg.training.warmup_epochs)
+
+        def lr_lambda(step):
+            if step < warmup_steps:
+                return step / warmup_steps
+            return max(0.0, (total_steps - step) / max(1, total_steps - warmup_steps))
+
+        self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
+        return self
+
+    def _optimizer_step(self) -> None:
+        if self.cfg.training.grad_clip > 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.training.grad_clip)
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.scheduler.step()
+        self.optimizer.zero_grad(set_to_none=True)
+
+    def train_one_epoch(self, epoch: int) -> float:
+        self.model.train()
+        criterion = nn.CrossEntropyLoss()
+        total_loss = 0.0
+        num_batches = 0
+        accum = self.cfg.training.grad_accumulation_steps
+        self.optimizer.zero_grad(set_to_none=True)
+        for batch_idx, batch in enumerate(self.train_loader):
+            x, meta, orig_Ts, labels, affines = prepare_batch_data(batch, self.device)
+            with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp):
+                logits = self.model(x, meta=meta, orig_Ts=orig_Ts, affines=affines)
+                loss = criterion(logits, labels)
+            self.scaler.scale(loss / accum).backward()
+            if (batch_idx + 1) % accum == 0:
+                self._optimizer_step()
+            total_loss += float(loss.item())
+            num_batches += 1
+            if self.rank == 0 and (batch_idx + 1) % self.cfg.logging.log_interval == 0:
+                self.logger.info("Epoch %d [%d/%d] loss=%.6f", epoch + 1, batch_idx + 1, len(self.train_loader), loss.item())
+        if num_batches % accum != 0:
+            self._optimizer_step()
+        return total_loss / max(1, num_batches)
+
+    @torch.no_grad()
+    def evaluate(self, loader, split_name: str) -> Dict[str, float]:
+        self.model.eval()
+        criterion = nn.CrossEntropyLoss()
+        preds, labels_all = [], []
+        total_loss = 0.0
+        num_batches = 0
+        for batch in loader:
+            x, meta, orig_Ts, labels, affines = prepare_batch_data(batch, self.device)
+            with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp):
+                logits = self.model(x, meta=meta, orig_Ts=orig_Ts, affines=affines)
+                loss = criterion(logits, labels)
+            total_loss += float(loss.item())
+            num_batches += 1
+            preds.extend(torch.argmax(logits, dim=1).cpu().numpy())
+            labels_all.extend(labels.cpu().numpy())
+        metrics = {
+            "loss": total_loss / max(1, num_batches),
+            "accuracy": accuracy_score(labels_all, preds),
+            "precision_macro": precision_score(labels_all, preds, average="macro", zero_division=0),
+            "recall_macro": recall_score(labels_all, preds, average="macro", zero_division=0),
+            "f1_macro": f1_score(labels_all, preds, average="macro", zero_division=0),
+            "f1_weighted": f1_score(labels_all, preds, average="weighted", zero_division=0),
+        }
+        if self.rank == 0:
+            self.logger.info("%s metrics: %s", split_name, metrics)
+        return metrics
+
+    def save(self, epoch: int, metrics: dict, is_best: bool):
+        if self.rank != 0:
+            return
+        os.makedirs(self.cfg.logging.checkpoint_dir, exist_ok=True)
+        payload = {
+            "epoch": epoch,
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "scheduler": self.scheduler.state_dict(),
+            "metrics": metrics,
+            "config": vars(self.cfg.model),
+        }
+        torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "downstream_latest.pt"))
+        if is_best:
+            torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "downstream_best.pt"))
+
+    def _load_best_for_test(self) -> None:
+        best_path = os.path.join(self.cfg.logging.checkpoint_dir, "downstream_best.pt")
+        if not os.path.exists(best_path):
+            return
+        checkpoint = torch.load(best_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint["model"])
+
+    def _save_test_metrics(self, metrics: Dict[str, float]) -> None:
+        if self.rank != 0:
+            return
+        os.makedirs(self.cfg.logging.checkpoint_dir, exist_ok=True)
+        path = os.path.join(self.cfg.logging.checkpoint_dir, "test_metrics.json")
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(metrics, f, indent=2)
+
+    def fit(self):
+        self.build()
+        if self.rank == 0:
+            self.logger.info("Starting downstream on %s", self.device)
+            self.logger.info("Train size=%d Val size=%d", len(self.train_loader.dataset), len(self.val_loader.dataset))
+        best_f1 = -1.0
+        for epoch in range(self.cfg.training.epochs):
+            train_loss = self.train_one_epoch(epoch)
+            val_metrics = self.evaluate(self.val_loader, "Validation")
+            metrics = {"val": val_metrics, "train_loss": train_loss}
+            is_best = val_metrics["f1_macro"] > best_f1
+            if is_best:
+                best_f1 = val_metrics["f1_macro"]
+            self.save(epoch, metrics, is_best=is_best)
+            if self.rank == 0:
+                self.logger.info("Epoch %d done train=%.6f best_f1=%.6f", epoch + 1, train_loss, best_f1)
+        if self.test_loader is not None:
+            self._load_best_for_test()
+            test_metrics = self.evaluate(self.test_loader, "Test")
+            self._save_test_metrics(test_metrics)

flexibrain/engine/pretrainer.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.optim as optim
+
+from flexibrain.config import RunConfig
+from flexibrain.data import build_pretrain_dataloaders
+from flexibrain.data.collate import prepare_batch_data
+from flexibrain.distributed import cleanup_distributed, setup_distributed
+from flexibrain.models import build_pretrain_model
+from flexibrain.utils.logging import setup_logger
+from flexibrain.utils.seed import set_seed
+from flexibrain.utils.training import get_dynamic_momentum, update_ema
+
+
+class Pretrainer:
+    def __init__(self, cfg: RunConfig):
+        self.cfg = cfg
+        self.rank = cfg.training.local_rank
+        self.world_size = cfg.training.world_size
+        if self.world_size > 1:
+            setup_distributed(self.rank, self.world_size)
+        self.device = torch.device(f"cuda:{self.rank}" if torch.cuda.is_available() else "cpu")
+        self.logger = setup_logger("pretrain", cfg.logging.log_dir, rank=self.rank)
+
+    def build(self):
+        set_seed(self.cfg.training.seed)
+        self.model = build_pretrain_model(self.cfg.model, self.device)
+        self.train_loader, self.val_loader = build_pretrain_dataloaders(self.cfg.data, self.cfg.training, rank=self.rank, world_size=self.world_size)
+        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.cfg.training.lr, weight_decay=self.cfg.training.weight_decay)
+        self.use_amp = bool(self.cfg.training.use_amp and self.device.type == "cuda")
+        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)
+        total_steps = max(1, len(self.train_loader) * self.cfg.training.epochs)
+        warmup_steps = max(1, len(self.train_loader) * self.cfg.training.warmup_epochs)
+
+        def lr_lambda(step):
+            if step < warmup_steps:
+                return step / warmup_steps
+            progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+            cycle_progress = (progress * 4) % 1.0
+            if cycle_progress < 0.8:
+                return 0.5 * (1 + math.cos(math.pi * cycle_progress / 0.8))
+            return 0.0
+
+        self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
+        return self
+
+    def _optimizer_step(self, momentum: float) -> None:
+        if self.cfg.training.grad_clip > 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.training.grad_clip)
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=True)
+        update_ema(self.model, momentum)
+        self.scheduler.step()
+
+    def train_one_epoch(self, epoch: int) -> float:
+        self.model.train()
+        total_loss = 0.0
+        num_batches = 0
+        accumulation_steps = self.cfg.training.grad_accumulation_steps
+        momentum = get_dynamic_momentum(epoch, self.cfg.training.epochs, self.cfg.model.momentum, self.cfg.model.final_momentum)
+        self.optimizer.zero_grad(set_to_none=True)
+        for batch_idx, batch in enumerate(self.train_loader):
+            x, meta, orig_Ts, affines = prepare_batch_data(batch, self.device)
+            with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp):
+                loss, _, _, _ = self.model(x, mask_ratio=self.cfg.training.mask_ratio, meta=meta, orig_Ts=orig_Ts, affines=affines)
+            self.scaler.scale(loss / accumulation_steps).backward()
+            total_loss += float(loss.item())
+            num_batches += 1
+            if (batch_idx + 1) % accumulation_steps == 0:
+                self._optimizer_step(momentum)
+            if self.rank == 0 and (batch_idx + 1) % self.cfg.logging.log_interval == 0:
+                self.logger.info("Epoch %d [%d/%d] loss=%.6f avg=%.6f momentum=%.6f", epoch + 1, batch_idx + 1, len(self.train_loader), loss.item(), total_loss / num_batches, momentum)
+        if num_batches % accumulation_steps != 0:
+            self._optimizer_step(momentum)
+        return total_loss / max(1, num_batches)
+
+    @torch.no_grad()
+    def validate(self, epoch: int) -> float:
+        self.model.eval()
+        total_loss = 0.0
+        num_batches = 0
+        for batch in self.val_loader:
+            x, meta, orig_Ts, affines = prepare_batch_data(batch, self.device)
+            with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp):
+                loss, _, _, _ = self.model(x, mask_ratio=self.cfg.training.mask_ratio, meta=meta, orig_Ts=orig_Ts, affines=affines)
+            total_loss += float(loss.item())
+            num_batches += 1
+        avg = total_loss / max(1, num_batches)
+        if self.rank == 0:
+            self.logger.info("Epoch %d validation loss=%.6f", epoch + 1, avg)
+        return avg
+
+    def save(self, epoch: int, val_loss: float, best_loss: float, is_best: bool):
+        if self.rank != 0:
+            return
+        os.makedirs(self.cfg.logging.checkpoint_dir, exist_ok=True)
+        payload = {
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+            "scheduler_state_dict": self.scheduler.state_dict(),
+            "val_loss": val_loss,
+            "best_loss": best_loss,
+            "config": vars(self.cfg.model),
+        }
+        torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "checkpoint_latest.pt"))
+        if is_best:
+            torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "checkpoint_best.pt"))
+
+    def fit(self):
+        self.build()
+        if self.rank == 0:
+            self.logger.info("Starting pretrain on %s", self.device)
+            self.logger.info("Train size=%d Val size=%d", len(self.train_loader.dataset), len(self.val_loader.dataset))
+        best_loss = float("inf")
+        for epoch in range(self.cfg.training.epochs):
+            if hasattr(self.train_loader.sampler, "set_epoch"):
+                self.train_loader.sampler.set_epoch(epoch)
+            train_loss = self.train_one_epoch(epoch)
+            val_loss = self.validate(epoch)
+            is_best = val_loss < best_loss
+            if is_best:
+                best_loss = val_loss
+            self.save(epoch, val_loss, best_loss, is_best=is_best)
+            if self.rank == 0:
+                self.logger.info("Epoch %d done train=%.6f val=%.6f best=%.6f", epoch + 1, train_loss, val_loss, best_loss)
+        if self.world_size > 1:
+            cleanup_distributed()

flexibrain/models/__init__.py

@@ -0,0 +1,4 @@
+from flexibrain.models.factory import build_downstream_model, build_pretrain_model
+from flexibrain.models.mamba_jepa import VolumeMambaJEPA
+
+__all__ = [build_downstream_model, build_pretrain_model, VolumeMambaJEPA]

flexibrain/models/classifier.py

@@ -0,0 +1,237 @@
+import torch
+import torch.nn as nn
+from flexibrain.models.transformer_block import Block
+
+from flexibrain.models.mamba_jepa import VolumeMambaJEPA
+
+
+class MambaJEPAClassifier(nn.Module):
+
+    def __init__(
+        self,
+        backbone: 'VolumeMambaJEPA',
+        num_classes: int,
+        head_depth: int = 2,
+        head_num_heads: int = 8,
+        head_mlp_ratio: float = 4.0,
+        head_qkv_bias: bool = True,
+        head_attn_drop: float = 0.0,
+        head_proj_drop: float = 0.0,
+        head_drop_path: float = 0.0,
+        head_norm_epsilon: float = 1e-5,
+        mlp_hidden: int = 1024,
+        mlp_depth: int = 2,
+        mlp_dropout: float = 0.1,
+        freeze_backbone: bool = False,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.embed_dim = backbone.embed_dim
+
+        if freeze_backbone:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+        else:
+            for p in self.backbone.parameters():
+                p.requires_grad = True
+
+        if dtype is None:
+            dtype = next(self.backbone.parameters()).dtype
+        if device is None:
+            device = next(self.backbone.parameters()).device
+        factory = dict(device=device, dtype=dtype)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim, **factory))
+        nn.init.normal_(self.cls_token, std=0.02)
+
+        dpr = [x.item() for x in torch.linspace(0, head_drop_path, head_depth)] if head_depth > 0 else []
+        def _norm_layer_with_dtype(dim):
+            ln = nn.LayerNorm(dim, eps=head_norm_epsilon)
+            if device is not None:
+                ln = ln.to(device=device)
+            if dtype is not None:
+                ln = ln.to(dtype=dtype)
+            return ln
+        self.head_blocks = nn.ModuleList([
+            Block(
+                dim=self.embed_dim,
+                num_heads=head_num_heads,
+                mlp_ratio=head_mlp_ratio,
+                qkv_bias=head_qkv_bias,
+                attn_drop=head_attn_drop,
+                drop=head_proj_drop,
+                drop_path=dpr[i] if head_depth > 0 else 0.0,
+                norm_layer=_norm_layer_with_dtype,
+            ) for i in range(head_depth)
+        ])
+        self.head_norm = nn.LayerNorm(self.embed_dim, eps=head_norm_epsilon, **factory)
+
+        layers = []
+        in_dim = self.embed_dim
+        for _ in range(max(mlp_depth - 1, 0)):
+            layers += [nn.Linear(in_dim, mlp_hidden, **factory), nn.GELU(), nn.Dropout(mlp_dropout)]
+            in_dim = mlp_hidden
+        layers += [nn.Linear(in_dim, num_classes, **factory)]
+        self.classifier = nn.Sequential(*layers)
+
+    @torch.no_grad()
+    def _encode_backbone_nograd(self, x, meta=None, orig_Ts=None, affines=None, inference_params=None):
+        xf, attn_pad, _, _ = self.backbone.patch_embed(x, meta, orig_Ts, affines)
+        feat = self.backbone._run_blocks(
+            xf, attn_pad,
+            blocks=self.backbone.blocks,
+            norm_layer=self.backbone.norm_f,
+            inference_params=inference_params
+        )
+        return feat, attn_pad
+
+    def _encode_backbone(self, x, meta=None, orig_Ts=None, affines=None, inference_params=None):
+        xf, attn_pad, _, _ = self.backbone.patch_embed(x, meta, orig_Ts, affines)
+
+        feat = self.backbone._run_blocks(
+            xf, attn_pad,
+            blocks=self.backbone.blocks,
+            norm_layer=self.backbone.norm_f,
+            inference_params=inference_params
+        )
+
+        return feat, attn_pad
+
+    def forward_from_tokens(self, tokens, attn_pad, inference_params=None):
+
+        B, L, D = tokens.shape
+        device = tokens.device
+
+        cls_tok = self.cls_token.to(dtype=tokens.dtype).expand(B, -1, -1)  # [B,1,D]
+        x_cat = torch.cat([cls_tok, tokens], dim=1)                        # [B,1+L,D]
+
+        cls_pad = torch.zeros(B, 1, dtype=torch.bool, device=device)     # False=valid
+        attn_cat = torch.cat([cls_pad, attn_pad], dim=1)                 # [B,1+L]
+
+        attn_cat_for_flash = ~attn_cat  
+
+        h = x_cat
+        for blk in self.head_blocks:
+            h = blk(h, attention_mask=attn_cat_for_flash)
+        h = self.head_norm(h)
+
+        cls_feat = h[:, 0, :]                                            # [B,D]
+        logits = self.classifier(cls_feat)                                # [B,C]
+        return logits
+
+    def forward(self, x, meta=None, orig_Ts=None, affines=None, inference_params=None):
+
+        feat, attn_pad = self._encode_backbone(x, meta=meta, orig_Ts=orig_Ts, affines=affines,
+                                             inference_params=inference_params)
+        B, L, D = feat.shape
+        device = feat.device
+
+        cls_tok = self.cls_token.to(dtype=feat.dtype).expand(B, -1, -1)  # [B,1,D]
+        x_cat = torch.cat([cls_tok, feat], dim=1)                        # [B,1+L,D]
+
+        cls_pad = torch.zeros(B, 1, dtype=torch.bool, device=device)     # False=valid
+        attn_cat = torch.cat([cls_pad, attn_pad], dim=1)                 # [B,1+L]
+
+        attn_cat_for_flash = ~attn_cat  
+
+        h = x_cat
+        for blk in self.head_blocks:
+            h = blk(h, attention_mask=attn_cat_for_flash)
+        h = self.head_norm(h)
+
+        cls_feat = h[:, 0, :]                                            # [B,D]
+        logits = self.classifier(cls_feat)                                # [B,C]
+        return logits
+
+class MambaJEPAClassifierAvgPool(nn.Module):
+
+    def __init__(
+        self,
+        backbone: 'VolumeMambaJEPA',
+        num_classes: int,
+        mlp_hidden: int = 1024,
+        mlp_depth: int = 3,
+        mlp_dropout: float = 0.1,
+        freeze_backbone: bool = False,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.embed_dim = backbone.embed_dim
+
+        if freeze_backbone:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+        else:
+            for p in self.backbone.parameters():
+                p.requires_grad = True
+
+        if dtype is None:
+            dtype = next(self.backbone.parameters()).dtype
+        if device is None:
+            device = next(self.backbone.parameters()).device
+        factory = dict(device=device, dtype=dtype)
+
+        layers = []
+        in_dim = self.embed_dim
+        for _ in range(max(mlp_depth - 1, 0)):
+            layers += [nn.Linear(in_dim, mlp_hidden, **factory), nn.GELU(), nn.Dropout(mlp_dropout)]
+            in_dim = mlp_hidden
+        layers += [nn.Linear(in_dim, num_classes, **factory)]
+        self.classifier = nn.Sequential(*layers)
+
+    def _encode_backbone(self, x, meta=None, orig_Ts=None, affines=None, inference_params=None):
+        xf, attn_pad, _, _ = self.backbone.patch_embed(x, meta, orig_Ts, affines)
+
+        feat = self.backbone._run_blocks(
+            xf, attn_pad,
+            blocks=self.backbone.blocks,
+            norm_layer=self.backbone.norm_f,
+            inference_params=inference_params
+        )
+        return feat, attn_pad
+
+    def forward_from_tokens(self, tokens, attn_pad, inference_params=None):
+        """
+        Args:
+            tokens: [B, L, D] backbone token
+            attn_pad: [B, L] attention mask (True=padding, False=valid)
+            inference_params
+
+        Returns:
+            logits: [B, num_classes]
+        """
+        B, L, D = tokens.shape
+
+        valid_mask = ~attn_pad  # [B, L] -> True=valid
+
+        valid_counts = valid_mask.sum(dim=1, keepdim=True).clamp(min=1)  # [B, 1]
+
+        feat_masked = tokens * valid_mask.unsqueeze(-1).float()  # [B, L, D]
+        feat_sum = feat_masked.sum(dim=1)  # [B, D]
+
+        feat_avg = feat_sum / valid_counts.float()  # [B, D]
+
+        logits = self.classifier(feat_avg)  # [B, num_classes]
+        return logits
+
+    def forward(self, x, meta=None, orig_Ts=None, affines=None, inference_params=None):
+
+        feat, attn_pad = self._encode_backbone(x, meta=meta, orig_Ts=orig_Ts, affines=affines,
+                                             inference_params=inference_params)
+        B, L, D = feat.shape
+
+        valid_mask = ~attn_pad  # [B, L] -> True=valid
+
+        valid_counts = valid_mask.sum(dim=1, keepdim=True).clamp(min=1)  # [B, 1]
+
+        feat_masked = feat * valid_mask.unsqueeze(-1).float()  # [B, L, D]
+        feat_sum = feat_masked.sum(dim=1)  # [B, D]
+
+        feat_avg = feat_sum / valid_counts.float()  # [B, D]
+
+        logits = self.classifier(feat_avg)  # [B, num_classes]
+        return logits

flexibrain/models/factory.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from flexibrain.config import ModelConfig, apply_checkpoint_config
+from flexibrain.models.mamba_jepa import VolumeMambaJEPA
+from flexibrain.models.classifier import MambaJEPAClassifier, MambaJEPAClassifierAvgPool
+
+
+def build_mamba_backbone(cfg: ModelConfig, device: torch.device, dtype=torch.float32) -> VolumeMambaJEPA:
+    return VolumeMambaJEPA(
+        embed_dim=cfg.embed_dim,
+        depth=cfg.depth,
+        predictor_depth=cfg.predictor_depth,
+        ssm_cfg=None,
+        encoder_attn_layer_idx=None,
+        attn_cfg=None,
+        drop_path_rate=cfg.drop_path_rate,
+        norm_epsilon=1e-5,
+        rms_norm=cfg.rms_norm,
+        initializer_cfg=None,
+        fused_add_norm=cfg.fused_add_norm,
+        residual_in_fp32=cfg.residual_in_fp32,
+        device=device,
+        dtype=dtype,
+        bimamba_type=cfg.bimamba_type,
+        if_bimamba=cfg.if_bimamba,
+        mixer_type=cfg.mixer_type,
+        if_devide_out=cfg.if_devide_out,
+        momentum=cfg.momentum,
+        norm_target=cfg.norm_target,
+    )
+
+
+def build_pretrain_model(cfg: ModelConfig, device: torch.device) -> nn.Module:
+    if cfg.model_type != "mamba":
+        raise ValueError("This cleaned Flexibrain build currently keeps only the Mamba pretrain/downstream path")
+    return build_mamba_backbone(cfg, device=device, dtype=torch.float32).to(device)
+
+
+def load_checkpoint(path: str, device: torch.device):
+    return torch.load(path, map_location=device)
+
+
+def state_dict_from_checkpoint(checkpoint: dict):
+    if "model_state_dict" in checkpoint:
+        return checkpoint["model_state_dict"]
+    if "model" in checkpoint:
+        return checkpoint["model"]
+    raise KeyError("Checkpoint has neither model_state_dict nor model")
+
+
+def build_downstream_model(cfg: ModelConfig, device: torch.device, logger: Optional[logging.Logger] = None, checkpoint_path: Optional[str] = None, from_scratch: bool = False, use_checkpoint_config: bool = True) -> nn.Module:
+    checkpoint = None
+    if checkpoint_path and not from_scratch:
+        checkpoint = load_checkpoint(checkpoint_path, device)
+        if use_checkpoint_config:
+            apply_checkpoint_config(cfg, checkpoint.get("config", {}))
+            if logger:
+                logger.info("Backbone config restored from checkpoint: %s", checkpoint.get("config", {}))
+    if cfg.model_type != "mamba":
+        raise ValueError("This cleaned Flexibrain build currently keeps only the Mamba downstream path")
+    backbone = build_mamba_backbone(cfg, device=device, dtype=torch.float32)
+    if checkpoint is not None:
+        state = state_dict_from_checkpoint(checkpoint)
+        try:
+            backbone.load_state_dict(state, strict=True)
+            if logger:
+                logger.info("Loaded pretrained backbone strictly from %s", checkpoint_path)
+        except RuntimeError as exc:
+            incompatible = backbone.load_state_dict(state, strict=False)
+            backward_markers = ["_b", "conv1d_b", "x_proj_b", "dt_proj_b", "A_b_log", "D_b"]
+            missing = list(incompatible.missing_keys)
+            only_backward = missing and all(any(marker in key for marker in backward_markers) for key in missing)
+            if not only_backward or incompatible.unexpected_keys:
+                raise exc
+            if logger:
+                logger.warning("Strict load missed %d backward-scan BiMamba keys; loaded checkpoint with strict=False compatibility", len(missing))
+    elif logger:
+        logger.info("Backbone initialized from scratch")
+    if cfg.head_type == "transformer":
+        model = MambaJEPAClassifier(
+            backbone=backbone,
+            num_classes=cfg.num_classes,
+            head_depth=cfg.head_depth,
+            head_num_heads=cfg.head_num_heads,
+            head_mlp_ratio=cfg.head_mlp_ratio,
+            head_proj_drop=cfg.head_proj_drop,
+            head_drop_path=cfg.head_drop_path,
+            mlp_hidden=cfg.mlp_hidden,
+            mlp_depth=cfg.mlp_depth,
+            mlp_dropout=cfg.mlp_dropout,
+            freeze_backbone=cfg.freeze_backbone,
+            device=device,
+        )
+    elif cfg.head_type == "avgpool":
+        model = MambaJEPAClassifierAvgPool(backbone=backbone, num_classes=cfg.num_classes, mlp_hidden=cfg.mlp_hidden, mlp_depth=cfg.mlp_depth, mlp_dropout=cfg.mlp_dropout, freeze_backbone=cfg.freeze_backbone, device=device)
+    else:
+        raise ValueError(f"Unknown head_type: {cfg.head_type}")
+    return model.to(device)

flexibrain/models/layers/moe.py

@@ -0,0 +1,156 @@
+import torch
+import torch.nn as nn
+
+
+class CondLoRA(nn.Module):
+    def __init__(self, dim, cond_in, rank=8, cond_hidden=32, device=None, dtype=None):
+        super().__init__()
+        fk={}
+        if device is not None: fk["device"]=device
+        if dtype  is not None: fk["dtype"]=dtype
+        self.U = nn.Parameter(torch.randn(dim, rank, **fk)*0.02)
+        self.V = nn.Parameter(torch.randn(dim, rank, **fk)*0.02)
+        self.cond = nn.Sequential(
+            nn.LayerNorm(cond_in),
+            nn.Linear(cond_in, cond_hidden, **fk), nn.GELU(),
+            nn.Linear(cond_hidden, rank, **fk)
+        )
+
+    def forward(self, x, cond_vec, attn_mask=None):   # x:[B,L,D], cond_vec:[B,C]
+        B,L,D = x.shape
+        a = self.cond(cond_vec)                        # [B,r]
+        xU = torch.einsum('bld,dr->blr', x, self.U)    # [B,L,r]
+        add = torch.einsum('blr,br,dr->bld', xU, a, self.V)  # [B,L,D]
+        y = x + add
+        if attn_mask is not None:
+            valid = (~attn_mask).unsqueeze(-1)
+            y = torch.where(valid, y, x)
+        return y
+
+class ResoPrior_E(nn.Module):
+    def __init__(self, k_space=0.5, k_time=0.5, gamma=0.5):
+        super().__init__()
+        self.ks, self.kt, self.g = k_space, k_time, gamma
+    def forward(self, rxyztr):
+        rx, ry, rz, tr = torch.clamp(rxyztr, 1e-6).unbind(-1)
+        sig = torch.stack([self.ks*rx, self.ks*ry, self.ks*rz, self.kt*tr], dim=-1)  # [B,4]
+        snr = (rx*ry*rz) / (tr ** self.g + 1e-6)                                     # [B]
+        z = torch.cat([torch.log(sig+1e-6), torch.log(snr+1e-6).unsqueeze(-1)], dim=-1)  # [B,5]
+        return z
+
+
+class ExpertMLP(nn.Module):
+    def __init__(self, dim, hidden_dim=None, device=None, dtype=None):
+        super().__init__()
+        hidden_dim = hidden_dim or dim * 4
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.fc1 = nn.Linear(dim, hidden_dim, **factory_kwargs)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, dim, **factory_kwargs)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+class MoE(nn.Module):
+    def __init__(self,
+                 dim,
+                 hidden_dim=None,   
+                 num_indep=3,
+                 aux_loss_coef=0.00,
+                 device=None,
+                 dtype=None,
+                 load_balance_coef: float = 0.01,
+                 use_res_cond: bool = False,
+                 cond_dim: int = 5,
+                 cond_hidden_dim: int = 16,
+                 cond_tanh_scale: float = 0.5):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.num_shared = 1
+        self.num_indep = int(num_indep)
+        self.num_experts = self.num_shared + self.num_indep
+        self.aux_loss_coef = float(aux_loss_coef)
+        self.load_balance_coef = load_balance_coef  
+
+        experts = [ExpertMLP(dim, hidden_dim, **factory_kwargs)]
+        for _ in range(self.num_indep):
+            experts.append(ExpertMLP(dim, hidden_dim, **factory_kwargs))
+        self.experts = nn.ModuleList(experts)
+
+        self.router_token = nn.Linear(dim, self.num_experts, bias=False, **factory_kwargs)
+
+        self.use_res_cond = bool(use_res_cond)
+        self.cond_tanh_scale = float(cond_tanh_scale)
+        if self.use_res_cond:
+            self.cond_proj = nn.Sequential(
+                nn.LayerNorm(cond_dim, **factory_kwargs),
+                nn.Linear(cond_dim, cond_hidden_dim, **factory_kwargs),
+                nn.GELU(),
+                nn.LayerNorm(cond_hidden_dim, **factory_kwargs),
+            )
+            self.router_scale = nn.Linear(cond_hidden_dim, self.num_experts, bias=False, **factory_kwargs)
+            self.router_bias  = nn.Linear(cond_hidden_dim, self.num_experts, bias=False, **factory_kwargs)
+        else:
+            self.router_scale = None
+            self.router_bias  = None
+
+        self.use_router_film = False          
+        if self.use_router_film and self.use_res_cond:
+            self.film_gamma = nn.Linear(cond_hidden_dim, dim, **factory_kwargs)
+            self.film_beta  = nn.Linear(cond_hidden_dim, dim, **factory_kwargs)
+
+
+    def forward(self, x, attn_mask=None, cond_vec: torch.Tensor = None, return_gates: bool = False):
+        """
+        x: [B, L, D]
+        attn_mask: [B, L]  True=pad, False=valid
+        cond_venc: [B, 3]
+        return_gates
+        return: y: [B, L, D], aux: scalar, (gates: [B, L, E] if return_gates=True)
+        """
+        B, L, D = x.shape
+
+        if self.use_res_cond:
+            cond = self.cond_proj(cond_vec)  # [B,cond_dim]
+            if self.use_router_film:
+                gamma = torch.tanh(self.film_gamma(cond))  # [B,D]
+                beta  = self.film_beta(cond)
+                x = x * (1 + 0.3 * gamma.unsqueeze(1)) + 0.3 * beta.unsqueeze(1)
+
+
+        token_logits = self.router_token(x)  # [B, L, E]
+
+        if self.use_res_cond:
+            scale = torch.tanh(self.router_scale(cond))           # [B,E]
+            bias  = self.router_bias(cond)                        # [B,E]
+            token_logits = token_logits * (1 + self.cond_tanh_scale * scale.unsqueeze(1)) \
+                                         + bias.unsqueeze(1)      # [B,L,E]
+
+        gates = torch.softmax(token_logits, dim=-1)             # [B, L, E]
+
+        if attn_mask is not None:
+            valid = ~attn_mask                            # [B, L]
+            gates = gates * valid.unsqueeze(-1)
+            gates = gates / gates.sum(dim=-1, keepdim=True).clamp_min(1e-6)
+
+        expert_outs = torch.stack([expert(x) for expert in self.experts], dim=-2)  # [B, L, E, D]
+        y = (gates.unsqueeze(-1) * expert_outs).sum(dim=-2)                        # [B, L, D]
+
+
+        imp = gates.sum(dim=(0, 1))                      # [E]
+        imp = imp / imp.sum().clamp_min(1e-6)           
+        uniform = torch.full_like(imp, 1.0 / self.num_experts)  
+        load_balance_loss = ((imp - uniform) ** 2).sum() * self.load_balance_coef  
+
+        aux = x.new_zeros(())
+        if self.aux_loss_coef > 0.0:
+            imp = gates.sum(dim=(0, 1))                      # [E]
+            imp = imp / imp.sum().clamp_min(1e-6)
+            uniform = torch.full_like(imp, 1.0 / self.num_experts)
+            aux = ((imp - uniform) ** 2).sum() * self.aux_loss_coef
+
+        if return_gates:
+            return y, load_balance_loss, gates
+        return y, load_balance_loss

flexibrain/models/layers/pos_embed.py

@@ -0,0 +1,62 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+def stape_patch_world_coords_physical(
+    X:int, Y:int, Z:int,
+    kx:int, ky:int, kz:int,
+    affine: torch.Tensor,
+    rho_mm: Tuple[float, float, float], 
+    device=None, dtype=None
+):
+    if device is None: device = affine.device
+    if dtype is None:  dtype = torch.float32
+
+    A = affine[:3, :3].to(device=device, dtype=dtype)  # [3,3]
+    t = affine[:3, 3].to(device=device, dtype=dtype)   # [3]
+
+    Lx, Ly, Lz = X//kx, Y//ky, Z//kz
+    icx = torch.arange(Lx, device=device, dtype=dtype)*kx + (kx-1)*0.5
+    icy = torch.arange(Ly, device=device, dtype=dtype)*ky + (ky-1)*0.5
+    icz = torch.arange(Lz, device=device, dtype=dtype)*kz + (kz-1)*0.5
+
+    gx, gy, gz = torch.meshgrid(icx, icy, icz, indexing='ij')     # [Lx,Ly,Lz]
+    idx = torch.stack([gx, gy, gz], dim=-1).reshape(-1, 3)        # [N,3]
+    coords = idx @ A.T + t                                        # [N,3]
+    return coords
+
+
+
+class FixedSinCos3DPE(nn.Module):
+    def __init__(self, embed_dim:int, num_freq:int=12,
+                 space_scale:float=1.0, learnable_proj: bool=True):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_freq = num_freq
+        freq = torch.exp(torch.linspace(0, math.log(10000.0), num_freq)) / 10000.0  
+        self.register_buffer('freq', freq)  # [num_freq]
+        self.space_scale = space_scale
+        in_dim = 3 * 2 * num_freq
+        self.proj = nn.Linear(in_dim, embed_dim, bias=False) if learnable_proj else nn.Identity()
+
+    def forward(self, xyz: torch.Tensor):
+        """
+        xyz: [B, L, 3]
+        return: [B, L, embed_dim]
+        """
+        B, L, _ = xyz.shape
+        x = xyz[..., 0] * self.space_scale
+        y = xyz[..., 1] * self.space_scale
+        z = xyz[..., 2] * self.space_scale
+
+        freq = self.freq.to(device=xyz.device, dtype=xyz.dtype)
+
+        def enc(u):
+            u = u[..., None] * freq   # [B,L,num_freq]
+            return torch.cat([torch.sin(u), torch.cos(u)], dim=-1)  # [B,L,2*num_freq]
+
+        feats = torch.cat([enc(x), enc(y), enc(z)], dim=-1)  # [B,L, 3*2*num_freq]
+        return self.proj(feats)  # [B,L,C]

flexibrain/models/layers/stape.py

@@ -0,0 +1,290 @@
+# stape_time_to_space.py
+import math
+from collections import defaultdict
+from typing import Dict, List, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from flexibrain.models.layers.pos_embed import FixedSinCos3DPE, stape_patch_world_coords_physical
+from flexibrain.utils.weight_resize import pi_resize_weight_1d, pi_resize_weight_3d
+
+class STAPE4D_TimeToSpace(nn.Module):
+    """
+    input:
+        x: (B, 96, 96, 96, T_max) 
+        meta: {subject_idx: {"voxel": (vx,vy,vz) mm, "tr": float s}}
+        orig_T
+        affine
+    output:
+        tokens:   (B, L_max, D_out)
+        attn_mask:(B, L_max)  True=padding
+        lengths:  List[int] 
+    """
+    def __init__(self,
+                 d_mid: int = 128,     
+                 d_out: int = 256,
+                 kt_base: int = 6,
+                 kx_base: int = 6,
+                 ky_base: int = 6,
+                 kz_base: int = 6,
+                 tau_seconds: float = 6.0,
+                 rho_mm: Tuple[float, float, float] = (12., 12., 12.),
+                 ):
+        super().__init__()
+        self.Dm = d_mid
+        self.Do = d_out
+        self.kt0, self.kx0, self.ky0, self.kz0 = kt_base, kx_base, ky_base, kz_base
+        self.tau = float(tau_seconds)
+        self.rho = tuple(float(r) for r in rho_mm)
+
+        # time based kernal [Dm, 1, kt0] 
+        self.w_t_first_base = nn.Parameter(
+            torch.randn(d_mid, 1, kt_base) * (1.0 / (1 * kt_base)) ** 0.5
+        )
+        self.b_t_first = nn.Parameter(torch.zeros(d_mid))
+
+        # space base kernal: [Do, Dm, kx0, ky0, kz0] 
+        self.w_xyz_after_base = nn.Parameter(
+            torch.randn(d_out, d_mid, kx_base, ky_base, kz_base) *
+            (1.0 / (d_mid * kx_base * ky_base * kz_base)) ** 0.5
+        )
+        self.b_xyz_after = nn.Parameter(torch.zeros(d_out))
+
+        self._cache_t = {}     # key=(kt,dtype) -> w_t_first
+        self._cache_xyz = {}   # key=(kx,ky,kz,dtype) -> w_xyz_after
+
+        # physical pe
+        self.pos_embed = FixedSinCos3DPE(
+            embed_dim=d_out,           
+            num_freq=12,               
+            space_scale=0.01,          
+            learnable_proj=True        
+        )
+
+    @torch.no_grad()
+    def _k_from_meta(self, tr: float, voxel: Tuple[float, float, float]) -> Tuple[int, int, int, int]:
+        vx, vy, vz = voxel
+        if tr <= 0 or not math.isfinite(float(tr)):
+            raise ValueError(f"TR must be positive for STAPE kernel sizing, got {tr!r}")
+        if any(v <= 0 or not math.isfinite(float(v)) for v in (vx, vy, vz)):
+            raise ValueError(f"Voxel spacing must be positive for STAPE kernel sizing, got {voxel!r}")
+        kt = max(1, round(self.tau / tr))
+        kx = max(1, round(self.rho[0] / vx))
+        ky = max(1, round(self.rho[1] / vy))
+        kz = max(1, round(self.rho[2] / vz))
+        return int(kt), int(kx), int(ky), int(kz)
+
+    def _get_wt_first(self, kt:int, device, dtype):
+        wt = pi_resize_weight_1d(self.w_t_first_base.to(dtype), kt)  # [Dm,1,kt]
+        return wt.to(device)
+
+    def _get_wxyz_after(self, kx:int, ky:int, kz:int, device, dtype):
+        w = pi_resize_weight_3d(self.w_xyz_after_base.to(dtype), kx, ky, kz)  # [Do,Dm,kx,ky,kz]
+        return w.to(device)
+
+    @staticmethod
+    def _detect_true_T(x_b: torch.Tensor) -> int:
+
+        with torch.no_grad():
+            s = x_b.abs().sum(dim=(0,1,2))
+            nz = torch.nonzero(s > 0, as_tuple=False)
+            if nz.numel() == 0:
+                return 0
+            return int(nz.max().item() + 1)
+
+    @staticmethod
+    def _spatial_keep_mask_alltime(x_b: torch.Tensor, kx:int, ky:int, kz:int, T_true:int) -> torch.Tensor:
+        X=Y=Z=96
+        Lx, Ly, Lz = X//kx, Y//ky, Z//kz
+        if T_true == 0:
+            return torch.zeros(Lx*Ly*Lz, dtype=torch.bool, device=x_b.device)
+        vol = (x_b[:,:,:,:T_true] != 0).any(dim=-1).float()  # [96,96,96] -> 1/0
+        vol = vol[:Lx*kx, :Ly*ky, :Lz*kz].unsqueeze(0).unsqueeze(0)  # [1,1,X,Y,Z]
+        keep = F.max_pool3d(vol, kernel_size=(kx,ky,kz), stride=(kx,ky,kz)) > 0  # [1,1,Lx,Ly,Lz]
+        return keep.squeeze(0).squeeze(0).reshape(-1)
+
+    def _compute_spatial_coords_for_group(self,
+                                        group_idxs: List[int],
+                                        affines: List[torch.Tensor],
+                                        kx: int, ky: int, kz: int,
+                                        device: torch.device) -> torch.Tensor:
+        """
+        compute patch physical coordinate
+        """
+        G = len(group_idxs)
+        X, Y, Z = 96, 96, 96  
+        Lx, Ly, Lz = X//kx, Y//ky, Z//kz
+
+        coords_list = []
+        for g, affine in enumerate(affines):
+            coords = stape_patch_world_coords_physical(
+                X=X, Y=Y, Z=Z,
+                kx=kx, ky=ky, kz=kz,
+                affine=affine,
+                rho_mm=self.rho  
+            )  # [Lx*Ly*Lz, 3]
+            coords_list.append(coords)
+
+        # [G, Lx*Ly*Lz, 3]
+        coords_batch = torch.stack(coords_list, dim=0)
+        return coords_batch.to(device)
+
+    def _add_positional_encoding(self,
+                               tokens_all: torch.Tensor,    # [G, Lx*Ly*Lz, Do]
+                               group_idxs: List[int],
+                               affines: List[torch.Tensor],
+                               kx: int, ky: int, kz: int) -> torch.Tensor:
+        device = tokens_all.device
+
+        coords = self._compute_spatial_coords_for_group(
+            group_idxs, affines, kx, ky, kz, device
+        )  # [G, Lx*Ly*Lz, 3]
+
+        pos_encoding = self.pos_embed(coords)  # [G, Lx*Ly*Lz, Do]  
+        pos_encoding = pos_encoding.to(tokens_all.dtype)
+        tokens_with_pos = tokens_all + pos_encoding
+
+        return tokens_with_pos, pos_encoding
+
+    def _run_group_time_first(self,
+                              x_group: torch.Tensor,            # [G,96,96,96,T_max]
+                              orig_Ts: List[int],               
+                              kt:int, kx:int, ky:int, kz:int,
+                              group_idxs: List[int],          
+                              affines: List[torch.Tensor],    
+                              return_grid_info: bool = False) -> Tuple[List[torch.Tensor], List[int], Dict]:
+        device, dtype = x_group.device, x_group.dtype
+        G, X, Y, Z, T_max = x_group.shape
+        assert X==96 and Y==96 and Z==96
+
+        T_true_max = max(orig_Ts) if len(orig_Ts)>0 else 0
+        T_pad = math.ceil(T_true_max / kt) * kt
+        T_prime = T_pad // kt   
+
+        w_t = self._get_wt_first(kt, device, dtype)   # [Dm,1,kt]
+
+        xg = x_group.clone()
+        if T_max < T_pad:
+            pad_len = T_pad - T_max
+            xg = F.pad(xg, (0, pad_len), mode='constant', value=0.0)  # [G,96,96,96,T_pad]
+        xg = xg[..., :T_pad]
+
+        xlin = xg.permute(0,1,2,3,4).contiguous().view(G*X*Y*Z, 1, T_pad)
+        b_t_first = self.b_t_first.to(device=device, dtype=dtype)
+        tfeat = F.conv1d(xlin, w_t, bias=b_t_first, stride=kt)  # [N, Dm, T′]
+        tfeat = tfeat.view(G, X, Y, Z, self.Dm, T_prime).permute(0,4,5,1,2,3).contiguous()
+        x_sp_in = tfeat.view(G, self.Dm*T_prime, X, Y, Z)  # [G, C_in, X,Y,Z]
+
+        w_xyz = self._get_wxyz_after(kx,ky,kz, device, dtype)   # [Do, Dm, kx,ky,kz]
+        w_xyz_rep = w_xyz.repeat(1, T_prime, 1, 1, 1)           # [Do, Dm*T′, kx,ky,kz]
+        b_xyz_after = self.b_xyz_after.to(device=device, dtype=dtype)
+        sfeat = F.conv3d(x_sp_in, w_xyz_rep, bias=b_xyz_after, stride=(kx,ky,kz))  # [G, Do, Lx,Ly,Lz]
+
+        Lx, Ly, Lz = X//kx, Y//ky, Z//kz
+        tokens_all = sfeat.permute(0,2,3,4,1).contiguous().view(G, Lx*Ly*Lz, self.Do)
+
+        if affines is not None and len(affines) == len(group_idxs):
+            tokens_all, pos_group = self._add_positional_encoding(tokens_all, group_idxs, affines, kx, ky, kz)
+
+        tokens_list, lengths, pos_list = [], [], []
+        grid_data = {}
+
+        for g in range(G):
+            T_true = orig_Ts[g]
+
+            keep_mask = self._spatial_keep_mask_alltime(x_group[g], kx,ky,kz, T_true)  # [Lx*Ly*Lz]
+            if keep_mask.any():
+                toks = tokens_all[g][keep_mask]   # [N_valid, Do]
+                pe = pos_group[g][keep_mask]
+            else:
+                toks = tokens_all[g].new_zeros((0, self.Do))
+                pe = pos_group[g].new_zeros((0, self.Do))
+
+            tokens_list.append(toks)
+            pos_list.append(pe)
+            lengths.append(int(toks.size(0)))
+
+            if return_grid_info:
+                sample_idx = group_idxs[g]
+                grid_data[sample_idx] = {
+                    'Lx': Lx,
+                    'Ly': Ly,
+                    'Lz': Lz,
+                    'kx': kx,
+                    'ky': ky,
+                    'kz': kz,
+                    'keep_mask': keep_mask.cpu(),  # [Lx*Ly*Lz] bool
+                    'grid_to_token_idx': torch.nonzero(keep_mask, as_tuple=False).squeeze(-1).cpu(),  
+                }
+
+        return tokens_list, lengths, grid_data, pos_list
+    
+    def forward(self,
+                x: torch.Tensor,                     
+                meta: Dict[int, Dict],               
+                orig_Ts: Sequence[int] = None,
+                affines: Sequence[torch.Tensor] = None,
+                return_grid_info: bool = False):         
+        B = x.size(0)
+        device, dtype = x.device, x.dtype
+
+        if orig_Ts is None:
+            orig_Ts = [self._detect_true_T(x[b]) for b in range(B)]
+        else:
+            orig_Ts = [int(t) for t in orig_Ts]
+
+        if affines is None:
+            print("WARNING: not provide affine")
+            affines = [torch.eye(4, device=device, dtype=dtype) for _ in range(B)]
+        else:
+            affines = [aff.to(device=device, dtype=dtype) for aff in affines]
+
+        groups = defaultdict(list)
+        for i in range(B):
+            voxel = tuple(meta[i]["voxel"])
+            tr = float(meta[i]["tr"])
+            group_key = (voxel, tr)
+            groups[group_key].append(i)
+
+        per_sample_tokens: List[torch.Tensor] = [None]*B
+        per_sample_pos:    List[torch.Tensor] = [None]*B
+        lengths: List[int] = [0]*B
+        grid_info: Dict[int, Dict] = {}  
+
+        for (voxel, tr), idxs in groups.items():
+            kt, kx, ky, kz = self._k_from_meta(tr, voxel)
+
+            x_group = x[idxs, ...]                              # [G,96,96,96,T_max]
+            Ts_group = [orig_Ts[i] for i in idxs]               
+            affines_group = [affines[i] for i in idxs]          
+
+            toks, lens, grid_data, pos_list = self._run_group_time_first(
+                x_group, Ts_group, kt, kx, ky, kz, idxs, affines_group,
+                return_grid_info=return_grid_info
+            )
+            for g_idx, (tok, ln) in enumerate(zip(toks, lens)):
+                loc = idxs[g_idx]
+                per_sample_tokens[loc] = tok
+                lengths[loc] = ln
+                per_sample_pos[loc]    = pos_list[g_idx]
+                if return_grid_info and loc in grid_data:
+                    grid_info[loc] = grid_data[loc]
+
+        L_max = max(lengths) if lengths else 0
+        out = x.new_zeros((B, L_max, self.Do))
+        pos_out    = x.new_zeros((B, L_max, self.Do))
+        attn_mask = torch.ones((B, L_max), dtype=torch.bool, device=device)
+
+        for b, tok in enumerate(per_sample_tokens):
+            n = lengths[b]
+            if n > 0:
+                out[b, :n] = tok
+                pos_out[b, :n] = per_sample_pos[b]
+                attn_mask[b, :n] = False
+
+        if return_grid_info:
+            return out, attn_mask, lengths, grid_info, pos_out
+        else:
+            return out, attn_mask, lengths, pos_out

flexibrain/models/mamba_blocks.py

@@ -0,0 +1,115 @@
+"""Mamba sequence blocks used by Flexibrain.
+
+References:
+- Mamba: https://github.com/state-spaces/mamba
+- 3D Mamba MAE: https://github.com/ydchen0806/TokenUnify
+
+This file keeps only the block factory pieces needed by the Flexibrain
+Mamba-JEPA backbone, instead of vendoring the full upstream training project.
+"""
+
+from functools import partial
+import inspect
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from timm.models.layers import DropPath
+
+from mamba_ssm.modules.mamba_simple import Mamba
+from mamba_ssm.modules.mamba2 import Mamba2
+from mamba_ssm.modules.mha import MHA
+from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn
+
+
+class Block(nn.Module):
+    def __init__(self, dim, mixer_cls, mlp_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False, drop_path=0.0):
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.norm = norm_cls(dim)
+        self.mixer = mixer_cls(dim)
+        try:
+            self._mixer_kwset = set(inspect.signature(self.mixer.forward).parameters.keys())
+        except Exception:
+            self._mixer_kwset = set()
+        if mlp_cls is not nn.Identity:
+            self.norm2 = norm_cls(dim)
+            self.mlp = mlp_cls(dim)
+        else:
+            self.mlp = None
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import failed"
+            assert isinstance(self.norm, (nn.LayerNorm, RMSNorm))
+
+    def forward(self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None, **mixer_kwargs):
+        if not self.fused_add_norm:
+            residual = (self.drop_path(hidden_states) + residual) if residual is not None else hidden_states
+            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+        else:
+            hidden_states, residual = layer_norm_fn(
+                self.drop_path(hidden_states), self.norm.weight, self.norm.bias,
+                residual=residual, prenorm=True, residual_in_fp32=self.residual_in_fp32,
+                eps=self.norm.eps, is_rms_norm=isinstance(self.norm, RMSNorm),
+            )
+        filtered_kwargs = {k: v for k, v in mixer_kwargs.items() if k in self._mixer_kwset}
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params, **filtered_kwargs)
+        if self.mlp is not None:
+            if not self.fused_add_norm:
+                residual = self.drop_path(hidden_states) + residual
+                residual = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
+                if self.residual_in_fp32:
+                    residual = residual.to(torch.float32)
+            else:
+                hidden_states, residual = layer_norm_fn(
+                    self.drop_path(hidden_states), self.norm2.weight, self.norm2.bias,
+                    residual=residual, prenorm=True, residual_in_fp32=self.residual_in_fp32,
+                    eps=self.norm2.eps, is_rms_norm=isinstance(self.norm2, RMSNorm),
+                )
+            hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+
+def create_block(d_model, ssm_cfg=None, attn_layer_idx=None, attn_cfg=None, norm_epsilon=1e-5, drop_path=0.0, rms_norm=False, residual_in_fp32=False, fused_add_norm=False, layer_idx=None, device=None, dtype=None, if_bimamba=False, bimamba_type="none", if_devide_out=False, init_layer_scale=None, mixer_type="mamba"):
+    if if_bimamba and bimamba_type == "none":
+        bimamba_type = "v1"
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    if attn_cfg is None:
+        attn_cfg = {}
+    factory_kwargs = {"device": device, "dtype": dtype}
+    if (attn_layer_idx is None) or (layer_idx not in attn_layer_idx):
+        if mixer_type == "mamba":
+            mixer_cls = partial(Mamba, layer_idx=layer_idx, init_layer_scale=init_layer_scale, bimamba_type=bimamba_type, if_devide_out=if_devide_out, **ssm_cfg, **factory_kwargs)
+        elif mixer_type == "mamba2":
+            mixer_cls = partial(Mamba2, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+        else:
+            raise ValueError(f"Unknown mixer_type: {mixer_type}")
+    else:
+        mixer_cls = partial(MHA, layer_idx=layer_idx, **attn_cfg, **factory_kwargs)
+    norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs)
+    block = Block(d_model, mixer_cls, nn.Identity, norm_cls=norm_cls, drop_path=drop_path, fused_add_norm=fused_add_norm, residual_in_fp32=residual_in_fp32)
+    block.layer_idx = layer_idx
+    return block
+
+
+def _init_weights(module, n_layer, initializer_range=0.02, rescale_prenorm_residual=True, n_residuals_per_layer=1):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None and not getattr(module.bias, "_no_reinit", False):
+            nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+    if rescale_prenorm_residual:
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)

flexibrain/models/mamba_jepa.py

@@ -0,0 +1,440 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+
+import copy
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from timm.models.vision_transformer import DropPath
+
+
+from flexibrain.models.mamba_blocks import RMSNorm, create_block, _init_weights, layer_norm_fn, rms_norm_fn
+
+
+from flexibrain.models.layers.stape import STAPE4D_TimeToSpace
+from flexibrain.models.layers.moe import MoE
+
+
+class VolumeMambaJEPA(nn.Module):
+    """ JEPA with VisionMamba backbone
+    """
+    def __init__(self,
+                embed_dim=512,
+                depth=24,
+                predictor_depth=2,
+                ssm_cfg=None, 
+                encoder_attn_layer_idx=None,
+                attn_cfg=None,
+                drop_path_rate=0.1,
+                norm_epsilon: float = 1e-5, 
+                rms_norm: bool = False, 
+                initializer_cfg=None,
+                fused_add_norm=True,
+                residual_in_fp32=True,
+                device=None,
+                dtype=None,
+                bimamba_type="none",
+                if_bimamba=False,
+                mixer_type="mamba",
+                if_devide_out=False,
+                momentum: float = 0.996,
+                norm_target: bool = True,
+
+                **kwargs
+                ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kwargs.update(factory_kwargs)
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.momentum = float(momentum)
+        self.norm_target = bool(norm_target)   
+
+        self.patch_embed = STAPE4D_TimeToSpace(
+            d_mid=16,
+            d_out=embed_dim,
+            kt_base=6,
+            kx_base=6,
+            ky_base=6,
+            kz_base=6,
+            tau_seconds=6.0,
+            rho_mm=(12.0, 12.0, 12.0),
+        )
+        if device is not None:
+            self.patch_embed = self.patch_embed.to(device=device)
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        inter_dpr = [0.0] + dpr
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        self.blocks = nn.ModuleList()
+        block_idx = 0
+        for i in range(depth):
+            self.blocks.append(
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    attn_layer_idx=encoder_attn_layer_idx,
+                    attn_cfg=attn_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=block_idx,
+                    bimamba_type=bimamba_type,
+                    if_bimamba=if_bimamba,
+                    drop_path=inter_dpr[i],
+                    if_devide_out=if_devide_out,   
+                    mixer_type=mixer_type,
+                    **factory_kwargs,
+                )
+            )
+            block_idx += 1
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            embed_dim, eps=norm_epsilon, **factory_kwargs
+        )
+
+        self.target_blocks = copy.deepcopy(self.blocks) 
+        self.target_norm = (nn.LayerNorm if not rms_norm else RMSNorm)(embed_dim, eps=norm_epsilon, **factory_kwargs)
+        for p in self.target_blocks.parameters():
+            p.requires_grad = False
+        for p in self.target_norm.parameters():
+            p.requires_grad = False
+
+
+        # ------- MoE & ResolutionxTR embed-------
+        self.moe_aux_coef = float(kwargs.get("moe_aux_coef", 0.1))
+        self.moe = MoE(
+            dim=embed_dim,
+            hidden_dim=embed_dim * 4,
+            num_indep=3,
+            aux_loss_coef=self.moe_aux_coef,
+            device=device,
+            dtype=dtype,
+        )
+
+        self.mask_token_ctx = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        torch.nn.init.normal_(self.mask_token_ctx, std=0.02)
+
+        self.pred_depth = predictor_depth
+        self.pred_dpr = [0.0 for _ in range(self.pred_depth)]
+        self.predictor_blocks = nn.ModuleList()
+        for i in range(self.pred_depth):
+            self.predictor_blocks.append(
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    attn_layer_idx=None,   
+                    attn_cfg=attn_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    bimamba_type=bimamba_type,
+                    if_bimamba=if_bimamba,
+                    drop_path=self.pred_dpr[i],
+                    if_devide_out=if_devide_out,
+                    mixer_type=mixer_type,
+                    **factory_kwargs,
+                )
+            )
+
+        self.predictor_norm = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            embed_dim, eps=norm_epsilon, **factory_kwargs)
+
+        self.apply(self._init_linear_ln)
+        self.initialize_jepa_weights()
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=depth,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+    def _init_linear_ln(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif (isinstance(m, nn.LayerNorm) or m.__class__.__name__ == 'RMSNorm') and hasattr(m, 'weight'):
+            nn.init.constant_(m.weight, 1.0)
+            if hasattr(m, 'bias') and m.bias is not None:
+                nn.init.constant_(m.bias, 0.0)
+
+    def initialize_jepa_weights(self):
+        torch.nn.init.normal_(self.mask_token_ctx, std=0.02)
+        if hasattr(self, 'predictor_blocks'):
+            self.predictor_blocks.apply(self._init_linear_ln)
+        if hasattr(self, 'predictor_norm'):
+            self._init_linear_ln(self.predictor_norm)
+        self.blocks.apply(self._init_linear_ln)
+        self._init_linear_ln(self.norm_f)
+
+        if hasattr(self.patch_embed, 'reset_parameters'):
+            self.patch_embed.reset_parameters()
+        else:
+            for m in self.patch_embed.modules():
+                if isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d):
+                    torch.nn.init.kaiming_normal_(m.weight, nonlinearity='linear')
+                    if m.bias is not None:
+                        nn.init.constant_(m.bias, 0)
+                elif isinstance(m, nn.Linear):
+                    torch.nn.init.xavier_uniform_(m.weight)
+                    if m.bias is not None:
+                        nn.init.constant_(m.bias, 0)
+
+        with torch.no_grad():
+            for ps, pt in zip(self.blocks.parameters(), self.target_blocks.parameters()):
+                pt.copy_(ps)
+            if hasattr(self, 'target_norm'):
+                for ps, pt in zip(self.norm_f.parameters(), self.target_norm.parameters()):
+                    pt.copy_(ps)
+    
+    @torch.no_grad()
+    def update_target_encoder(self, m: float = None):
+        """EMA update of the target encoder"""
+        m = float(m or self.momentum)
+        for ps, pt in zip(self.blocks.parameters(), self.target_blocks.parameters()):
+            pt.data.mul_(m).add_(ps.data, alpha=1 - m)
+        for ps, pt in zip(self.norm_f.parameters(), self.target_norm.parameters()):
+            pt.data.mul_(m).add_(ps.data, alpha=1 - m)
+    
+    
+    def random_masking(self, x, attn_mask, lengths, mask_ratio):
+        """
+        x: [B, Lmax, D], attn_mask: [B, Lmax] (True=pad), lengths: list[int]
+        Return:
+        x_keep: [B, Lk_max, D]
+        mask_full: [B, Lmax] (0=keep, 1=remove; pad 仍为1)
+        ids_restore: [B, Lmax]
+        attn_keep: [B, Lk_max] (True=pad)
+        keep_lengths: list[int]
+        ids_keep_pad: [B, Lk_max]  
+        """
+        N, Lmax, D = x.shape
+        device = x.device
+
+        x_keep_list, mask_list, ids_restore_list = [], [], []
+        ids_keep_list, keep_lengths, Lk_max = [], [], 0
+
+        for i in range(N):
+            Li = lengths[i]
+            if Li == 0:
+                x_keep_list.append(torch.empty(0, D, device=device, dtype=x.dtype))
+                mask_list.append(torch.ones(Lmax, device=device))
+                ids_restore_list.append(torch.arange(Lmax, device=device))
+                ids_keep_list.append(torch.empty(0, dtype=torch.long, device=device))
+                keep_lengths.append(0)
+                continue
+
+            Lk = max(1, int(Li * (1 - mask_ratio)))
+            keep_lengths.append(Lk)
+            Lk_max = max(Lk_max, Lk)
+
+            noise = torch.rand(Li, device=device)
+            ids_shuffle = torch.argsort(noise)           # [Li]
+            ids_restore_valid = torch.argsort(ids_shuffle)
+            ids_keep = ids_shuffle[:Lk]                 
+
+            x_keep_list.append(x[i, ids_keep])
+            ids_keep_list.append(ids_keep)
+
+            mask_i = torch.ones(Lmax, device=device)
+            valid_mask = torch.ones(Li, device=device)
+            valid_mask[:Lk] = 0
+            mask_i[:Li] = torch.gather(valid_mask, 0, ids_restore_valid)
+            mask_list.append(mask_i)
+
+            ids_restore_full = torch.arange(Lmax, device=device)
+            ids_restore_full[:Li] = ids_restore_valid
+            ids_restore_list.append(ids_restore_full)
+
+        # pad keep tensors
+        x_keep = x.new_zeros((N, max(1, Lk_max), D))
+        ids_keep_pad = torch.full((N, max(1, Lk_max)), -1, dtype=torch.long, device=device)
+        attn_keep = torch.ones(N, max(1, Lk_max), dtype=torch.bool, device=device)
+
+        for i, (xi, ik) in enumerate(zip(x_keep_list, ids_keep_list)):
+            if xi.numel() > 0:
+                Lk = xi.size(0)
+                x_keep[i, :Lk] = xi
+                ids_keep_pad[i, :Lk] = ik
+                attn_keep[i, :Lk] = False
+
+        mask_full = torch.stack(mask_list, dim=0)
+        ids_restore = torch.stack(ids_restore_list, dim=0)
+
+        return x_keep, mask_full, ids_restore, attn_keep, keep_lengths, ids_keep_pad
+
+    def _build_context_visible(self, x_keep, attn_keep):
+
+        return x_keep, attn_keep
+    
+    def _build_target_masked(self, x_full, mask_full, lengths):
+        """
+          x_tgt_pad:   [B, Lt_max, D]
+          attn_tgt:    [B, Lt_max] (True=padding)
+          tgt_lengths: list[int]
+        """
+        B, Lmax, D = x_full.shape
+        device, dtype = x_full.device, x_full.dtype
+        per_sample, tgt_lengths, Lt_max = [], [], 0
+        for i in range(B):
+            Li = lengths[i]
+            if Li == 0:
+                per_sample.append(x_full.new_empty((0, D)))
+                tgt_lengths.append(0)
+                continue
+            sel = (mask_full[i, :Li] == 1) if mask_full.dtype != torch.bool else mask_full[i, :Li]
+            xi = x_full[i, :Li][sel]
+            per_sample.append(xi)
+            tgt_lengths.append(xi.size(0))
+            Lt_max = max(Lt_max, xi.size(0))
+
+        x_tgt_pad = x_full.new_zeros((B, Lt_max, D))
+        attn_tgt = torch.ones(B, Lt_max, dtype=torch.bool, device=device)
+        for i, xi in enumerate(per_sample):
+            if xi.numel() > 0:
+                Lti = xi.size(0)
+                x_tgt_pad[i, :Lti] = xi
+                attn_tgt[i, :Lti] = False
+
+        return x_tgt_pad, attn_tgt, tgt_lengths
+
+    def _run_blocks(self, x, attn_mask, blocks, norm_layer, inference_params=None, unpack_buffer=None):
+        residual = None
+        # x, seq_idx, idx_info = pack_batch(x, attn_mask)
+        hidden_states = x
+        for layer in blocks:
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params,
+                attn_mask=attn_mask
+            )
+            # hidden_states, residual = layer(
+            #     hidden_states, residual, inference_params=inference_params,
+            #     seq_idx=seq_idx
+            # )
+        fused_norm_available = layer_norm_fn is not None and (
+            RMSNorm is None or not isinstance(norm_layer, RMSNorm) or rms_norm_fn is not None
+        )
+        if not self.fused_add_norm or not fused_norm_available:
+            if residual is None:
+                residual = hidden_states
+            else:
+                residual = residual + self.drop_path(hidden_states)
+            hidden_states = norm_layer(residual.to(dtype=norm_layer.weight.dtype))
+        else:
+            fused_add_norm_fn = rms_norm_fn if RMSNorm is not None and isinstance(norm_layer, RMSNorm) else layer_norm_fn
+            hidden_states = fused_add_norm_fn(
+                self.drop_path(hidden_states),
+                norm_layer.weight,
+                norm_layer.bias,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+                eps=norm_layer.eps,
+            )
+        # hidden_states = unpack_batch(hidden_states, idx_info, unpack_buffer)
+        return hidden_states
+
+    def forward(self, x, mask_ratio=0.6, meta=None, orig_Ts=None, affines=None, inference_params=None, return_moe_features=False):
+
+        x_full, attn_pad, lengths, _ = self.patch_embed(x, meta, orig_Ts, affines, return_grid_info=False)
+
+        # random mask
+        x_keep, mask_full, ids_restore, attn_keep, keep_lengths, ids_keep_pad = self.random_masking(x_full, attn_pad, lengths, mask_ratio)
+
+        # build ctx visible
+        x_ctx, attn_ctx = self._build_context_visible(x_keep, attn_keep)
+        ctx_keep_out = self._run_blocks(x_ctx, attn_ctx,
+            blocks=self.blocks,
+            norm_layer=self.norm_f,
+        ) # [B, Lk_max, D]
+
+        device = x_full.device
+        B, Lmax, D = x_full.shape
+        ctx_keep_out, moe_aux, gates = self.moe(ctx_keep_out, attn_mask=attn_ctx, cond_vec=None, return_gates=True)
+
+        if return_moe_features:
+            return ctx_keep_out, attn_keep, keep_lengths, meta
+
+        ctx_full = x_full.new_zeros((B, Lmax, D))
+        attn_full = torch.ones(B, Lmax, dtype=torch.bool, device=device)
+        for i in range(B):
+            Li = lengths[i]
+            if Li == 0:
+                continue
+
+            Lk_i = int((~attn_keep[i]).sum().item())
+            if Lk_i > 0:
+                keep_idx = ids_keep_pad[i, :Lk_i].long()          # [Lk_i]
+                ctx_full[i, keep_idx] = ctx_keep_out[i, :Lk_i]    # [Lk_i, D]
+
+            masked_sel = (mask_full[i, :Li] == 1) if mask_full.dtype != torch.bool else mask_full[i, :Li]
+            if masked_sel.any():
+                idx = torch.nonzero(masked_sel, as_tuple=False).squeeze(1)   # [n_mask]
+                token_rows = self.mask_token_ctx[0, 0].expand(idx.numel(), D)  # [n_mask, D]
+                ctx_full[i, :Li].index_copy_(0, idx, token_rows)
+
+            attn_full[i, :Li] = False
+
+        pred_full = self._run_blocks(
+            ctx_full, attn_full, 
+            blocks=self.predictor_blocks,
+            norm_layer=self.predictor_norm,
+        ) # [B, Lmax, D]
+
+        x_tgt_pad, attn_tgt, tgt_lengths = self._build_target_masked(x_full, mask_full, lengths)
+        with torch.no_grad():
+            tgt_feat = self._run_blocks(
+                x_tgt_pad, attn_tgt, 
+                blocks=self.target_blocks,
+                norm_layer=self.target_norm,
+            ) # [B, Lt_max, D]
+
+        Lt_max = x_tgt_pad.size(1)
+        pred_masked = pred_full.new_zeros((B, Lt_max, D))
+        attn_pred = torch.ones(B, Lt_max, dtype=torch.bool, device=device)
+        for i in range(B):
+            Li = lengths[i]
+            if Li == 0:
+                continue
+            sel = (mask_full[i, :Li] == 1) if mask_full.dtype != torch.bool else mask_full[i, :Li]
+            vi = pred_full[i, :Li][sel]
+            if vi.numel() > 0:
+                Lti = vi.size(0)
+                pred_masked[i, :Lti] = vi
+                attn_pred[i, :Lti] = False
+
+        if self.norm_target:
+            # def norm_sg(x, eps=1e-6):
+            #     return x / (x.norm(dim=-1, keepdim=True).clamp_min(eps).detach())
+            # tgt_feat   = norm_sg(tgt_feat)
+            # pred_masked = norm_sg(pred_masked)
+            tgt_norm = torch.linalg.norm(tgt_feat, dim=-1, keepdim=True).clamp_min(1e-6)
+            tgt_feat = tgt_feat / tgt_norm
+            pred_norm = torch.linalg.norm(pred_masked, dim=-1, keepdim=True).clamp_min(1e-6)
+            pred_masked = pred_masked / pred_norm
+
+
+        valid = ~attn_pred
+        denom = valid.sum().clamp_min(1)
+        loss = (pred_masked[valid] - tgt_feat[valid]).pow(2).sum() / denom
+
+        loss = loss
+
+
+        return loss, pred_masked, tgt_feat, mask_full

flexibrain/models/transformer_block.py

@@ -0,0 +1,92 @@
+"""Transformer block used by Flexibrain downstream heads.
+
+Reference:
+- Brain-Harmony / BrainHarmonix official codebase: https://github.com/hzlab/Brain-Harmony
+- The official README marks the project license as CC BY-NC-SA 4.0.
+
+Only the small Block/Attention/MLP subset required by the downstream head is
+kept here; the rest of the Brain-Harmony repository is intentionally not
+vendored into Flexibrain.
+"""
+
+import torch
+import torch.nn as nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()
+    return x.div(keep_prob) * random_tensor
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob=0.0):
+        super().__init__()
+        self.drop_prob = float(drop_prob)
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return self.drop(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attention_mask=None, output_attentions=False):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        if attention_mask is not None:
+            valid = attention_mask.bool()
+            key_mask = ~valid[:, None, None, :]
+            attn = attn.masked_fill(key_mask, torch.finfo(attn.dtype).min)
+        attn = self.attn_drop(attn.softmax(dim=-1))
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj_drop(self.proj(x))
+        return (x, attn) if output_attentions else (x, None)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_mode=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.norm2 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.mlp = MLP(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+
+    def forward(self, x, attention_mask=None, return_attention=False):
+        y, attn = self.attn(self.norm1(x), attention_mask=attention_mask, output_attentions=return_attention)
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return (x, attn) if return_attention else x

flexibrain/utils/checkpoint.py

@@ -0,0 +1,108 @@
+import torch.nn as nn
+import torch
+from typing import Tuple
+import torch.optim as optim
+from torch.nn.parallel import DistributedDataParallel as DDP
+import os
+
+def load_checkpoint(
+    model: nn.Module,
+    optimizer: optim.Optimizer,
+    scheduler,
+    checkpoint_path: str,
+    device: torch.device,
+) -> Tuple[int, float]:
+    """Load model checkpoint."""
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    # Load model state dict (handle DDP wrapper)
+    if isinstance(model, DDP):
+        model.module.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint['model_state_dict'])
+
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+    if scheduler is not None and 'scheduler_state_dict' in checkpoint:
+        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+
+    epoch = checkpoint.get('epoch', 0)
+    best_loss = checkpoint.get('best_loss', float('inf'))
+
+    return epoch, best_loss
+
+
+def save_checkpoint(
+    model: nn.Module,
+    optimizer: optim.Optimizer,
+    scheduler,
+    epoch: int,
+    best_loss: float,
+    checkpoint_dir: str,
+    args=None,
+    rank: int = 0,
+) -> None:
+    """Save model checkpoint with configuration."""
+    if rank != 0:
+        return
+
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    # Get model state dict (handle DDP wrapper)
+    model_state = model.module.state_dict() if isinstance(model, DDP) else model.state_dict()
+
+    checkpoint = {
+        'epoch': epoch,
+        'model_state_dict': model_state,
+        'optimizer_state_dict': optimizer.state_dict(),
+        'best_loss': best_loss,
+    }
+
+    if scheduler is not None:
+        checkpoint['scheduler_state_dict'] = scheduler.state_dict()
+
+    # Save model configuration for downstream tasks
+    if args is not None:
+        checkpoint['config'] = {
+            'model_type': args.model_type,
+            'embed_dim': args.embed_dim,
+            'depth': args.depth,
+            'predictor_depth': args.predictor_depth,
+            'drop_path_rate': args.drop_path_rate,
+            'rms_norm': args.rms_norm,
+            'fused_add_norm': args.fused_add_norm,
+            'residual_in_fp32': args.residual_in_fp32,
+            'bimamba_type': args.bimamba_type,
+            'if_bimamba': args.if_bimamba,
+            'mixer_type': args.mixer_type,
+            'if_devide_out': args.if_devide_out,
+            'predictor_hidden': args.predictor_hidden,
+            'momentum': args.momentum,
+            'norm_target': args.norm_target,
+            'num_heads': args.num_heads,
+            'mlp_ratio': args.mlp_ratio,
+        }
+
+    # Save latest checkpoint
+    latest_path = os.path.join(checkpoint_dir, 'checkpoint_latest.pt')
+    torch.save(checkpoint, latest_path)
+
+    # Save best checkpoint
+    if best_loss is not None:
+        best_path = os.path.join(checkpoint_dir, 'checkpoint_best.pt')
+        torch.save(checkpoint, best_path)
+
+
+def save_downstream_checkpoint(model, optimizer, scheduler, epoch, metrics, checkpoint_dir, rank=0):
+    """Save downstream checkpoint."""
+    if rank == 0:
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        checkpoint = {
+            'epoch': epoch,
+            'model': model.state_dict() if not isinstance(model, DDP) else model.module.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'scheduler': scheduler.state_dict(),
+            'metrics': metrics,
+        }
+        path = os.path.join(checkpoint_dir, f"downstream_epoch_{epoch:03d}.pt")
+        torch.save(checkpoint, path)

flexibrain/utils/logging.py

@@ -0,0 +1,19 @@
+import logging
+import os
+from datetime import datetime
+
+
+def setup_logger(name: str, log_dir: str, rank: int = 0) -> logging.Logger:
+    os.makedirs(log_dir, exist_ok=True)
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO if rank == 0 else logging.WARNING)
+    logger.handlers.clear()
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    if rank == 0:
+        file_handler = logging.FileHandler(os.path.join(log_dir, f"{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"))
+        file_handler.setFormatter(formatter)
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+        logger.addHandler(console_handler)
+    return logger

flexibrain/utils/pinv_resize.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025
+# Utilities to build resize operators and their pseudoinverses.
+
+from typing import Tuple
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from functools import lru_cache
+
+@torch.no_grad()
+def _resize_2d(x: Tensor, shape: Tuple[int, int],
+               interpolation: str = "bicubic",
+               antialias: bool = True) -> Tensor:
+    """
+    Resize a 2D tensor x[h0,w0] -> shape[h,w] using torch interpolate.
+    Matches the "wrap with [None,None,...]" trick from your flex_patch_embed.py.
+    """
+    x_resized = F.interpolate(
+        x[None, None, ...],
+        shape,
+        mode=interpolation,
+        antialias=antialias,
+    )
+    return x_resized[0, 0, ...]
+
+
+@lru_cache(maxsize=256)
+def _calculate_pinv_2d(old_shape: Tuple[int, int],
+                       new_shape: Tuple[int, int],
+                       interpolation: str = "bicubic",
+                       antialias: bool = True,
+                       device: torch.device = torch.device("cpu"),
+                       dtype: torch.dtype = torch.float32) -> Tensor:
+    """
+    Build the (flattened) resize matrix R s.t. vec(new) = R @ vec(old),
+    then return pinv(R). This mirrors your flex_patch_embed.py approach.
+
+    Args:
+        old_shape: (h0, w0)
+        new_shape: (h, w)
+    Returns:
+        pinv(R): Tensor of shape [(h*w), (h0*w0)]
+    """
+    # Construct R by sending basis vectors through the geometric resize op.
+    mat = []
+    h0, w0 = int(old_shape[0]), int(old_shape[1])
+    for i in range(int(np.prod(old_shape))):
+        basis = torch.zeros((h0, w0), dtype=dtype, device=device)
+        idx = np.unravel_index(i, (h0, w0))
+        basis[idx] = 1.0
+        mat.append(_resize_2d(basis, new_shape, interpolation, antialias).reshape(-1))
+    resize_matrix = torch.stack(mat)  # [(h*w), (h0*w0)]
+    pinv = torch.linalg.pinv(resize_matrix)
+    return pinv  # [(h*w), (h0*w0)]

flexibrain/utils/seed.py

@@ -0,0 +1,11 @@
+import random
+import numpy as np
+import torch
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)

flexibrain/utils/training.py

@@ -0,0 +1,38 @@
+import numpy as np
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+def meta_to_matrix(meta: dict, B: int) -> np.ndarray:
+    out = np.empty((B, 4), dtype=np.float32)
+    for i in range(B):
+
+        m = meta[i]
+        voxel = m.get("voxel", m.get("voxel_size", m.get("spacing")))
+
+        rx = float(voxel[0])
+        ry = float(voxel[1])
+        rt = float(m.get("rt", voxel[2]))
+        tr = float(m["tr"])  
+
+        out[i] = (rx, ry, rt, tr)
+    return out
+
+def update_ema(model: nn.Module, momentum: float) -> None:
+    """Update target encoder with EMA."""
+    if hasattr(model, 'update_target_encoder'):
+        model.update_target_encoder(m=momentum)
+    elif isinstance(model, DDP) and hasattr(model.module, 'update_target_encoder'):
+        model.module.update_target_encoder(m=momentum)
+
+
+def get_dynamic_momentum(epoch: int, total_epochs: int, base_momentum: float = 0.996, final_momentum: float = 0.9999) -> float:
+    """
+    Calculate dynamic momentum for EMA.
+
+    Momentum increases from base_momentum to final_momentum over training.
+    This helps stabilize training in later epochs.
+    """
+    progress = epoch / total_epochs
+    # Cosine annealing: start at base, end at final
+    momentum = final_momentum - (final_momentum - base_momentum) * 0.5 * (1 + np.cos(np.pi * progress))
+    return momentum

flexibrain/utils/weight_resize.py

@@ -0,0 +1,131 @@
+# Weight (kernel) resizing using pseudoinverse-based geometric operators.
+import torch
+from torch import Tensor
+from flexibrain.utils.pinv_resize import _calculate_pinv_2d
+
+# ------------- 1D (e.g., time) -----------------
+def resize_conv1d_weight_with_pinv(
+    w_star: Tensor, k_new: int,
+    interpolation: str = "bicubic",
+    antialias: bool = True,
+) -> Tensor:
+    """
+    Resample a Conv1d kernel from K_old -> k_new using pinv of a 2D operator
+    on a degenerate dimension (1,K). This keeps the math aligned with the 2D codepath.
+
+    Args:
+        w_star: [Out, In, K_old]
+        k_new:  new kernel length
+    Returns:
+        w_new:  [Out, In, k_new]
+    """
+    Out, In, K_old = w_star.shape
+    if k_new == K_old:
+        return w_star
+
+    dev, dt = w_star.device, w_star.dtype
+    requires_grad = w_star.requires_grad
+
+    # Build pinv((1,K_old)->(1,k_new)) - 这个操作不需要梯度
+    with torch.no_grad():
+        pinv = _calculate_pinv_2d(
+            (1, int(K_old)), (1, int(k_new)),
+            interpolation=interpolation,
+            antialias=antialias,
+            device=dev, dtype=dt
+        )  # [(1*k_new), (1*K_old)] == [k_new, K_old]
+
+    W = w_star.reshape(Out * In, K_old)      # [(Out*In), K_old]
+    W_new = (pinv @ W.T).T                   # [(Out*In), k_new]
+    W_new = W_new.reshape(Out, In, k_new)
+
+    # 恢复requires_grad状态
+    if requires_grad:
+        W_new = W_new.requires_grad_(True)
+
+    return W_new
+
+
+def pi_resize_weight_1d(
+    w_star: Tensor, k_new: int,
+    interpolation: str = "bicubic",
+    antialias: bool = True,
+) -> Tensor:
+    """
+    Alias kept for timetospace: same signature as your current helper.
+    """
+    return resize_conv1d_weight_with_pinv(
+        w_star, k_new, interpolation=interpolation, antialias=antialias
+    )
+
+# ------------- 3D separable (x,y,z) ------------
+def resize_conv3d_weight_separable_with_pinv(
+    w_star: Tensor, kx: int, ky: int, kz: int,
+    interpolation: str = "bicubic",
+    antialias: bool = True,
+) -> Tensor:
+    """
+    Separable 3-axis resize using three 1D pinv operators.
+    Mirrors your existing axis-by-axis path; only changes how we form each 1D pinv.
+
+    Args:
+        w_star: [Out, In, Kx0, Ky0, Kz0]
+        kx, ky, kz: target sizes
+    Returns:
+        w_new:  [Out, In, kx, ky, kz]
+    """
+    Out, In, Kx0, Ky0, Kz0 = w_star.shape
+    if (kx, ky, kz) == (Kx0, Ky0, Kz0):
+        return w_star
+
+    dev, dt = w_star.device, w_star.dtype
+    requires_grad = w_star.requires_grad
+    W = w_star
+
+    # x-axis
+    with torch.no_grad():
+        Rx_p = _calculate_pinv_2d(
+            (int(Kx0), 1), (int(kx), 1),
+            device=dev, dtype=dt,
+            interpolation=interpolation, antialias=antialias
+        )  # [kx, Kx0]
+    W = W.permute(2, 0, 1, 3, 4).reshape(Kx0, -1)           # [Kx0, Out*In*Ky0*Kz0]
+    W = (Rx_p @ W).reshape(kx, Out, In, Ky0, Kz0).permute(1, 2, 0, 3, 4)
+
+    # y-axis
+    with torch.no_grad():
+        Ry_p = _calculate_pinv_2d(
+            (int(Ky0), 1), (int(ky), 1),
+            device=dev, dtype=dt,
+            interpolation=interpolation, antialias=antialias
+        )  # [ky, Ky0]
+    W = W.permute(3, 0, 1, 2, 4).reshape(Ky0, -1)
+    W = (Ry_p @ W).reshape(ky, Out, In, kx, Kz0).permute(1, 2, 3, 0, 4)
+
+    # z-axis
+    with torch.no_grad():
+        Rz_p = _calculate_pinv_2d(
+            (int(Kz0), 1), (int(kz), 1),
+            device=dev, dtype=dt,
+            interpolation=interpolation, antialias=antialias
+        )  # [kz, Kz0]
+    W = W.permute(4, 0, 1, 2, 3).reshape(Kz0, -1)
+    W = (Rz_p @ W).reshape(kz, Out, In, kx, ky).permute(1, 2, 3, 4, 0)
+
+    # 恢复requires_grad状态
+    if requires_grad:
+        W = W.requires_grad_(True)
+
+    return W
+
+def pi_resize_weight_3d(
+    w_star: Tensor, kx: int, ky: int, kz: int,
+    interpolation: str = "bicubic",
+    antialias: bool = True,
+) -> Tensor:
+    """
+    Alias kept for timetospace: same signature as your current helper.
+    """
+    return resize_conv3d_weight_separable_with_pinv(
+        w_star, kx, ky, kz, interpolation=interpolation, antialias=antialias
+    )

licenses/causal_conv1d_LICENSE_BSD_3_Clause.txt

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

licenses/mamba2_LICENSE_Apache_2.0.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2023 Tri Dao, Albert Gu
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

licenses/mamba_mae_LICENSE_CC_BY_NC_4.0.txt

@@ -0,0 +1,400 @@
+
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+   wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More_considerations
+     for the public: 
+   wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

mamba_ssm/__init__.py

@@ -0,0 +1,6 @@
+__version__ = "2.0.3"
+
+from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn
+from .modules.mamba_simple import Mamba
+from .modules.mamba2 import Mamba2
+from .models.mixer_seq_simple import MambaLMHeadModel

mamba_ssm/distributed/distributed_utils.py

@@ -0,0 +1,144 @@
+from typing import Optional
+
+import torch
+from torch import Tensor
+from torch.distributed import ProcessGroup
+
+# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for
+# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent
+# version of PyTorch. The following 4 lines are for backward compatibility with
+# older PyTorch.
+if "all_gather_into_tensor" not in dir(torch.distributed):
+    torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base
+if "reduce_scatter_tensor" not in dir(torch.distributed):
+    torch.distributed.reduce_scatter_tensor = torch.distributed._reduce_scatter_base
+
+
+# Raw operation, does not support autograd, but does support async
+def all_gather_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
+    world_size = torch.distributed.get_world_size(process_group)
+    output = torch.empty(
+        world_size * input_.shape[0], *input_.shape[1:], dtype=input_.dtype, device=input_.device
+    )
+    handle = torch.distributed.all_gather_into_tensor(
+        output, input_.contiguous(), group=process_group, async_op=async_op
+    )
+    return output, handle
+
+
+# Raw operation, does not support autograd, but does support async
+def reduce_scatter_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
+    world_size = torch.distributed.get_world_size(process_group)
+    assert input_.shape[0] % world_size == 0
+    output = torch.empty(
+        input_.shape[0] // world_size, *input_.shape[1:], dtype=input_.dtype, device=input_.device
+    )
+    handle = torch.distributed.reduce_scatter_tensor(
+        output, input_.contiguous(), group=process_group, async_op=async_op
+    )
+    return output, handle
+
+
+# Raw operation, does not support autograd, but does support async
+def all_reduce_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
+    input_ = input_.contiguous()
+    handle = torch.distributed.all_reduce(input_, group=process_group, async_op=async_op)
+    return input_, handle
+
+
+class AllGatherFunc(torch.autograd.Function):
+    """Gather the input from sequence parallel region and concatenate."""
+
+    @staticmethod
+    def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
+        ctx.process_group = process_group
+        output, _ = all_gather_raw(input_, process_group)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        grad_input, _ = reduce_scatter_raw(grad_output, ctx.process_group)
+        return grad_input, None
+
+
+# Supports autograd, but does not support async
+all_gather = AllGatherFunc.apply
+
+
+class ReduceScatterFunc(torch.autograd.Function):
+    """Reduce scatter the input from the sequence parallel region and concatenate."""
+
+    @staticmethod
+    def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
+        ctx.process_group = process_group
+        output, _ = reduce_scatter_raw(input_, process_group)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        grad_input, _ = all_gather_raw(grad_output, ctx.process_group)
+        return grad_input, None
+
+
+# Supports autograd, but does not support async
+reduce_scatter = ReduceScatterFunc.apply
+
+
+class AllReduceFunc(torch.autograd.Function):
+    """Gather the input from sequence parallel region and concatenate."""
+
+    @staticmethod
+    def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
+        ctx.process_group = process_group
+        output, _ = all_reduce_raw(input_, process_group)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        return grad_output, None
+
+
+# Supports autograd, but does not support async
+all_reduce = AllReduceFunc.apply
+
+
+def sync_shared_params(model: torch.nn.Module, process_group: ProcessGroup):
+    # We want to iterate over parameters with _shared_params=True in the same order,
+    # as different ranks might have different number of parameters (e.g., only rank 0 has bias).
+    pamams_shared = {
+        name: p for name, p in model.named_parameters() if getattr(p, "_shared_params", False)
+    }
+    for _, p in sorted(pamams_shared.items()):
+        with torch.no_grad():
+            # Broadcast needs src to be global rank, not group rank
+            torch.distributed.broadcast(
+                p, src=torch.distributed.get_global_rank(process_group, 0), group=process_group
+            )
+
+
+# Ref: https://github.com/NVIDIA/Megatron-LM/blob/52e636888cccc41e931251c417a7181fc36de926/megatron/optimizer/optimizer.py#L256
+def allreduce_sequence_parallel_grad(model: torch.nn.Module, process_group: ProcessGroup):
+    # We want to iterate over parameters with _sequence_parallel=True in the same order,
+    # as different ranks might have different number of parameters (e.g., only rank 0 has bias).
+    params_seqparallel = {
+        name: p for name, p in model.named_parameters() if getattr(p, "_sequence_parallel", False)
+    }
+    grads = [p.grad for _, p in sorted(params_seqparallel.items())]
+    if grads:
+        with torch.no_grad():
+            coalesced = torch._utils._flatten_dense_tensors(grads)
+            torch.distributed.all_reduce(coalesced, group=process_group)
+            for buf, synced in zip(grads, torch._utils._unflatten_dense_tensors(coalesced, grads)):
+                buf.copy_(synced)
+
+
+def get_dim_for_local_rank(dim: int, world_size: int, local_rank: int, multiple_of: int = 1) -> int:
+    """Get the dim for the local rank derived from splitting dim on world_size processes.
+
+    The split may not be even across the world_size processes.
+    """
+    multiple = dim // multiple_of
+    div = multiple // world_size
+    mod = multiple % world_size
+    local_multiple = div + int(local_rank < mod)
+    return local_multiple * multiple_of

mamba_ssm/distributed/tensor_parallel.py

@@ -0,0 +1,296 @@
+# Copyright (c) 2024, Tri Dao.
+# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.distributed import ProcessGroup
+
+from einops import rearrange
+
+from .distributed_utils import (
+    all_gather_raw,
+    all_reduce,
+    all_reduce_raw,
+    reduce_scatter,
+    reduce_scatter_raw,
+)
+
+
+class ParallelLinearFunc(torch.autograd.Function):
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, x, weight, bias, process_group=None, sequence_parallel=True):
+        """
+        If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
+        with sequence parallelism: we do an all_gather_raw of x before doing the matmul.
+        """
+        ctx.compute_weight_gradient = weight.requires_grad
+        ctx.process_group = process_group
+        ctx.sequence_parallel = sequence_parallel
+
+        if torch.is_autocast_enabled():
+            x = x.to(dtype=torch.get_autocast_gpu_dtype())
+        x = x.contiguous()
+        if process_group is not None and sequence_parallel:
+            # We want to kick off the all_gather early, before weight dtype conversion
+            total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
+        else:
+            total_x = x
+
+        if torch.is_autocast_enabled():
+            weight = weight.to(dtype=torch.get_autocast_gpu_dtype())
+            bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None
+        weight = weight.contiguous()
+        if process_group is not None and sequence_parallel:
+            handle_x.wait()
+        batch_shape, n = total_x.shape[:-1], total_x.shape[-1]
+        batch_dim = batch_shape.numel()
+        # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174
+        output = F.linear(total_x, weight, bias)
+        if ctx.compute_weight_gradient:
+            ctx.save_for_backward(x, weight)
+        else:
+            ctx.save_for_backward(weight)
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad_output):
+        grad_output = grad_output.contiguous()
+        process_group = ctx.process_group
+        sequence_parallel = ctx.sequence_parallel
+        if ctx.compute_weight_gradient:
+            x, weight = ctx.saved_tensors
+            if process_group is not None and sequence_parallel:
+                total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
+            else:
+                total_x = x
+        else:
+            (weight,) = ctx.saved_tensors
+            total_x = None
+        batch_shape = grad_output.shape[:-1]
+        batch_dim = batch_shape.numel()
+        grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
+        if ctx.needs_input_grad[0]:
+            grad_input = F.linear(grad_output, weight.t())
+            grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1])
+            if process_group is not None:
+                reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw
+                grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True)
+        else:
+            grad_input = None
+        if ctx.needs_input_grad[1]:
+            assert ctx.compute_weight_gradient
+            if process_group is not None and sequence_parallel:
+                handle_x.wait()
+            grad_weight = torch.einsum(
+                "bo,bi->oi", grad_output, total_x.reshape(batch_dim, total_x.shape[-1])
+            )
+        else:
+            grad_weight = None
+        grad_bias = grad_output.sum(dim=0) if ctx.needs_input_grad[2] else None
+        if process_group is not None and ctx.needs_input_grad[0]:
+            handle_grad_input.wait()
+        return grad_input, grad_weight, grad_bias, None, None
+
+
+def parallel_linear_func(
+    x: Tensor,
+    weight: Tensor,
+    bias: Optional[Tensor] = None,
+    process_group: Optional[ProcessGroup] = None,
+    sequence_parallel: bool = True,
+):
+    return ParallelLinearFunc.apply(x, weight, bias, process_group, sequence_parallel)
+
+
+class ColumnParallelLinear(nn.Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        process_group: ProcessGroup,
+        bias: bool = True,
+        sequence_parallel=True,
+        multiple_of=1,
+        device=None,
+        dtype=None,
+    ) -> None:
+        world_size = torch.distributed.get_world_size(process_group)
+        if out_features % multiple_of:
+            raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}")
+        multiple = out_features // multiple_of
+        # We want to split @multiple across world_size, but it could be an uneven split
+        div = multiple // world_size
+        mod = multiple % world_size
+        # The first @mod ranks get @div + 1 copies, the rest get @div copies
+        local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
+        super().__init__(
+            in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype
+        )
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+
+    def forward(self, x):
+        # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism:
+        # we do an all_gather of x before doing the matmul.
+        # If not, then the input is already gathered.
+        return parallel_linear_func(
+            x,
+            self.weight,
+            self.bias,
+            process_group=self.process_group,
+            sequence_parallel=self.sequence_parallel,
+        )
+
+
+class RowParallelLinear(nn.Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        process_group: ProcessGroup,
+        bias: bool = True,
+        sequence_parallel=True,
+        multiple_of=1,
+        device=None,
+        dtype=None,
+    ) -> None:
+        world_size = torch.distributed.get_world_size(process_group)
+        rank = torch.distributed.get_rank(process_group)
+        if in_features % multiple_of:
+            raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}")
+        multiple = in_features // multiple_of
+        # We want to split @multiple across world_size, but it could be an uneven split
+        div = multiple // world_size
+        mod = multiple % world_size
+        # The first @mod ranks get @div + 1 copies, the rest get @div copies
+        local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
+        # Only rank 0 will have bias
+        super().__init__(
+            local_multiple * multiple_of,
+            out_features,
+            bias=bias and rank == 0,
+            device=device,
+            dtype=dtype,
+        )
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+
+    def forward(self, x):
+        """
+        We're doing Tensor Parallel with sequence parallelism: we do the matmul and then
+        a reduce_scatter of the result.
+        """
+        out = parallel_linear_func(x, self.weight, self.bias)
+        reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
+        return reduce_fn(out, self.process_group)
+
+
+class VocabParallelEmbedding(nn.Embedding):
+    def __init__(self, num_embeddings, *args, process_group=None, padding_idx=None, **kwargs):
+        self.process_group = process_group
+        if process_group is not None:
+            world_size = torch.distributed.get_world_size(process_group)
+            if num_embeddings % world_size != 0:
+                raise ValueError(
+                    f"num_embeddings ({num_embeddings}) must be divisible by "
+                    f"world_size ({world_size})"
+                )
+            if world_size > 1 and padding_idx is not None:
+                raise RuntimeError("ParallelEmbedding does not support padding_idx")
+        else:
+            world_size = 1
+        super().__init__(num_embeddings // world_size, *args, padding_idx=padding_idx, **kwargs)
+
+    def forward(self, input: Tensor) -> Tensor:
+        if self.process_group is None:
+            return super().forward(input)
+        else:
+            rank = torch.distributed.get_rank(self.process_group)
+            vocab_size = self.num_embeddings
+            vocab_start_index, vocab_end_index = rank * vocab_size, (rank + 1) * vocab_size
+            # Create a mask of valid vocab ids (1 means it needs to be masked).
+            input_ids_mask = (input < vocab_start_index) | (input >= vocab_end_index)
+            input = input - vocab_start_index
+            input[input_ids_mask] = 0
+            embeddings = super().forward(input)
+            embeddings[input_ids_mask] = 0.0
+            return embeddings
+
+
+class ColumnParallelEmbedding(nn.Embedding):
+    def __init__(self, num_embeddings, embedding_dim, *args, process_group=None, **kwargs):
+        self.process_group = process_group
+        if process_group is not None:
+            world_size = torch.distributed.get_world_size(process_group)
+            if embedding_dim % world_size != 0:
+                raise ValueError(
+                    f"embedding_dim ({embedding_dim}) must be divisible by "
+                    f"world_size ({world_size})"
+                )
+        else:
+            world_size = 1
+        super().__init__(num_embeddings, embedding_dim // world_size, *args, **kwargs)
+
+
+class ParallelEmbeddings(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        vocab_size,
+        max_position_embeddings,
+        process_group,
+        padding_idx=None,
+        sequence_parallel=True,
+        device=None,
+        dtype=None,
+    ):
+        """
+        If max_position_embeddings <= 0, there's no position embeddings
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+        self.word_embeddings = VocabParallelEmbedding(
+            vocab_size,
+            embed_dim,
+            padding_idx=padding_idx,
+            process_group=process_group,
+            **factory_kwargs,
+        )
+        self.max_position_embeddings = max_position_embeddings
+        if self.max_position_embeddings > 0:
+            self.position_embeddings = ColumnParallelEmbedding(
+                max_position_embeddings, embed_dim, process_group=process_group, **factory_kwargs
+            )
+
+    def forward(self, input_ids, position_ids=None, combine_batch_seqlen_dim=False):
+        """
+        input_ids: (batch, seqlen)
+        position_ids: (batch, seqlen)
+        """
+        batch_size, seqlen = input_ids.shape
+        world_size = torch.distributed.get_world_size(self.process_group)
+        embeddings = self.word_embeddings(input_ids)
+        if self.max_position_embeddings > 0:
+            if position_ids is None:
+                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
+            position_embeddings = self.position_embeddings(position_ids)
+            if world_size <= 1:
+                embeddings = embeddings + position_embeddings
+            else:
+                partition_dim = self.position_embeddings.embedding_dim
+                rank = torch.distributed.get_rank(self.process_group)
+                embeddings[
+                    ..., rank * partition_dim : (rank + 1) * partition_dim
+                ] += position_embeddings
+        if combine_batch_seqlen_dim:
+            embeddings = rearrange(embeddings, "b s d -> (b s) d")
+        reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
+        return embeddings if world_size <= 1 else reduce_fn(embeddings, self.process_group)

mamba_ssm/models/config_mamba.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass, field
+
+
+@dataclass
+class MambaConfig:
+
+    d_model: int = 2560
+    d_intermediate: int = 0
+    n_layer: int = 64
+    vocab_size: int = 50277
+    ssm_cfg: dict = field(default_factory=dict)
+    attn_layer_idx: list = field(default_factory=list)
+    attn_cfg: dict = field(default_factory=dict)
+    rms_norm: bool = True
+    residual_in_fp32: bool = True
+    fused_add_norm: bool = True
+    pad_vocab_size_multiple: int = 8
+    tie_embeddings: bool = True

mamba_ssm/models/mixer_seq_simple.py

@@ -0,0 +1,307 @@
+# Copyright (c) 2023, Albert Gu, Tri Dao.
+
+import math
+from functools import partial
+import json
+import os
+import copy
+
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+
+from .config_mamba import MambaConfig
+from ..modules.mamba_simple import Mamba
+from ..modules.mamba2 import Mamba2
+from ..modules.mha import MHA
+from ..modules.mlp import GatedMLP
+from ..modules.block import Block
+from ..utils.generation import GenerationMixin
+from ..utils.hf import load_config_hf, load_state_dict_hf
+
+from ..ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn
+
+
+
+def create_block(
+    d_model,
+    d_intermediate,
+    ssm_cfg=None,
+    attn_layer_idx=None,
+    attn_cfg=None,
+    norm_epsilon=1e-5,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=False,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+):
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    if attn_layer_idx is None:
+        attn_layer_idx = []
+    if attn_cfg is None:
+        attn_cfg = {}
+    factory_kwargs = {"device": device, "dtype": dtype}
+    if layer_idx not in attn_layer_idx:
+        # Create a copy of the config to modify
+        ssm_cfg = copy.deepcopy(ssm_cfg) if ssm_cfg is not None else {}
+        ssm_layer = ssm_cfg.pop("layer", "Mamba1")
+        if ssm_layer not in ["Mamba1", "Mamba2"]:
+            raise ValueError(f"Invalid ssm_layer: {ssm_layer}, only support Mamba1 and Mamba2")
+        mixer_cls = partial(
+            Mamba2 if ssm_layer == "Mamba2" else Mamba,
+            layer_idx=layer_idx,
+            **ssm_cfg,
+            **factory_kwargs
+        )
+    else:
+        mixer_cls = partial(MHA, layer_idx=layer_idx, **attn_cfg, **factory_kwargs)
+    norm_cls = partial(
+        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
+    )
+    if d_intermediate == 0:
+        mlp_cls = nn.Identity
+    else:
+        mlp_cls = partial(
+            GatedMLP, hidden_features=d_intermediate, out_features=d_model, **factory_kwargs
+        )
+    block = Block(
+        d_model,
+        mixer_cls,
+        mlp_cls,
+        norm_cls=norm_cls,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+class MixerModel(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_layer: int,
+        d_intermediate: int,
+        vocab_size: int,
+        ssm_cfg=None,
+        attn_layer_idx=None,
+        attn_cfg=None,
+        norm_epsilon: float = 1e-5,
+        rms_norm: bool = False,
+        initializer_cfg=None,
+        fused_add_norm=False,
+        residual_in_fp32=False,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)
+
+        # We change the order of residual and layer norm:
+        # Instead of LN -> Attn / MLP -> Add, we do:
+        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
+        # the main branch (output of MLP / Mixer). The model definition is unchanged.
+        # This is for performance reason: we can fuse add + layer_norm.
+        self.fused_add_norm = fused_add_norm
+        if self.fused_add_norm:
+            if layer_norm_fn is None or rms_norm_fn is None:
+                raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")
+
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    d_model,
+                    d_intermediate=d_intermediate,
+                    ssm_cfg=ssm_cfg,
+                    attn_layer_idx=attn_layer_idx,
+                    attn_cfg=attn_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    **factory_kwargs,
+                )
+                for i in range(n_layer)
+            ]
+        )
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            d_model, eps=norm_epsilon, **factory_kwargs
+        )
+
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+                n_residuals_per_layer=1 if d_intermediate == 0 else 2,  # 2 if we have MLP
+            )
+        )
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(self, input_ids, inference_params=None, **mixer_kwargs):
+        hidden_states = self.embedding(input_ids)
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+        if not self.fused_add_norm:
+            residual = (hidden_states + residual) if residual is not None else hidden_states
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm=False here since we don't need the residual
+            hidden_states = layer_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+                is_rms_norm=isinstance(self.norm_f, RMSNorm)
+            )
+        return hidden_states
+
+
+class MambaLMHeadModel(nn.Module, GenerationMixin):
+
+    def __init__(
+        self,
+        config: MambaConfig,
+        initializer_cfg=None,
+        device=None,
+        dtype=None,
+    ) -> None:
+        self.config = config
+        d_model = config.d_model
+        n_layer = config.n_layer
+        d_intermediate = config.d_intermediate
+        vocab_size = config.vocab_size
+        ssm_cfg = config.ssm_cfg
+        attn_layer_idx = config.attn_layer_idx
+        attn_cfg = config.attn_cfg
+        rms_norm = config.rms_norm
+        residual_in_fp32 = config.residual_in_fp32
+        fused_add_norm = config.fused_add_norm
+        pad_vocab_size_multiple = config.pad_vocab_size_multiple
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        super().__init__()
+        if vocab_size % pad_vocab_size_multiple != 0:
+            vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)
+        self.backbone = MixerModel(
+            d_model=d_model,
+            n_layer=n_layer,
+            d_intermediate=d_intermediate,
+            vocab_size=vocab_size,
+            ssm_cfg=ssm_cfg,
+            attn_layer_idx=attn_layer_idx,
+            attn_cfg=attn_cfg,
+            rms_norm=rms_norm,
+            initializer_cfg=initializer_cfg,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+            **factory_kwargs,
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
+
+        # Initialize weights and apply final processing
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+        self.tie_weights()
+
+    def tie_weights(self):
+        if self.config.tie_embeddings:
+            self.lm_head.weight = self.backbone.embedding.weight
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, **mixer_kwargs):
+        """
+        "position_ids" is just to be compatible with Transformer generation. We don't use it.
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+        hidden_states = self.backbone(input_ids, inference_params=inference_params, **mixer_kwargs)
+        if num_last_tokens > 0:
+            hidden_states = hidden_states[:, -num_last_tokens:]
+        lm_logits = self.lm_head(hidden_states)
+        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
+        return CausalLMOutput(logits=lm_logits)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):
+        config_data = load_config_hf(pretrained_model_name)
+        config = MambaConfig(**config_data)
+        model = cls(config, device=device, dtype=dtype, **kwargs)
+        model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))
+        return model
+
+    def save_pretrained(self, save_directory):
+        """
+        Minimal implementation of save_pretrained for MambaLMHeadModel.
+        Save the model and its configuration file to a directory.
+        """
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, 'pytorch_model.bin')
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, 'config.json')
+        with open(config_path, 'w') as f:
+            json.dump(self.config.__dict__, f, indent=4)