Validation Dispatch

dispatch

FUNCTION	DESCRIPTION
dispatch_cross_validation	Low-level framework for parallel K-fold cross-validation with multiple datasets.
dispatch_bootstrap	Low-level framework for parallel bootstrap resampling with multiple datasets.

Classes

Functions

dispatch_cross_validation(estimator: Any | dict[str, Any], datasets: dict[str, DataBatch], fit_score_fn: Callable, n_cv: int | Any = 5, shuffle: bool = False, random_state: int | None = None, device: str | torch.device | None = None, n_jobs: int = 1, backend: str = 'threading', verbose: int = 0, pre_dispatch: str = '2*n_jobs', return_required: list[str] = ['score_test', 'score_train'], return_estimator: bool = False, return_times: bool = True, error_score: float | str = 'raise', dtype: torch.dtype = torch.float32) -> dict[str, Any]

Low-level framework for parallel K-fold cross-validation with multiple datasets.

Handles dataset validation, fold generation, parallel execution, and result aggregation for custom cross-validation workflows. Users provide a callback function that defines fit/score logic for each fold.

This is the core dispatch function; for single-dataset CV, use cross_validate().

PARAMETER	DESCRIPTION
estimator	Estimator(s) to validate per fold: Single object: Copied and used for all datasets Dict: Maps dataset names to estimators (keys must match datasets) TYPE: Any | dict[str, Any]
datasets	Multi-dataset mapping with required constraints: Keys: Arbitrary names (e.g., 'X', 'Y', 'metadata') Values: DataBatch objects with data, weights, target Constraint: All must have identical sample count Example: {'X': DataBatch(X, X_w, None), 'Y': DataBatch(Y, Y_w, y_target)} TYPE: dict[str, DataBatch]
fit_score_fn	Callback with signature: fit_score_fn(estimator, train_data, test_data) -> dict[str, Tensor] train_data: dict mapping dataset names → DataBatch (train split) test_data: dict mapping dataset names → DataBatch (test split) Returns dict with keys matching return_required All returned tensors must be scalars TYPE: Callable
n_cv	CV strategy: int: Create KFold(n_splits=n_cv, ...) Splitter object: sklearn-compatible with get_n_splits() and split() TYPE: int | Any DEFAULT: 5
shuffle	Whether to shuffle before splitting. Only used when n_cv is int. TYPE: bool DEFAULT: False
random_state	Random seed for shuffle. Only used when n_cv is int. Default is None. TYPE: int | None DEFAULT: None
device	Device for computation. Default uses device of first dataset. CUDA device disables parallelization (n_jobs -> 1). TYPE: str | device | None DEFAULT: None
n_jobs	Parallel jobs. -1 uses all processors. Automatically set to 1 when using CUDA device. TYPE: int DEFAULT: 1
backend	Joblib backend ('threading', 'multiprocessing', 'loky', etc.). TYPE: str, by default "threading" DEFAULT: 'threading'
verbose	Verbosity level (0=silent, higher=more output). TYPE: int, by default 0 DEFAULT: 0
pre_dispatch	Pre-dispatch strategy for joblib queue. TYPE: str, by default "2*n_jobs" DEFAULT: '2*n_jobs'
return_required	Keys that fit_score_fn must return in its output dict. TYPE: list[str], by default ["score_test", "score_train"] DEFAULT: ['score_test', 'score_train']
return_estimator	Whether to include fitted fold estimators in output. TYPE: bool, by default False DEFAULT: False
return_times	Whether to include fold execution times in output. TYPE: bool DEFAULT: True
error_score	Behavior on fit/score errors: 'raise': Re-raise exception 'warn': Assign nan and issue warning numeric: Assign that value TYPE: float | str DEFAULT: "raise"
dtype	Data type for error score tensors. TYPE: dtype DEFAULT: torch.float32

RETURNS	DESCRIPTION
dict[str, Any]	Dictionary with keys: 'score_test': Test scores per fold, shape (n_cv,) 'score_train': Train scores per fold, shape (n_cv,) (only if in return_required) 'time': Fold times in seconds, length n_cv (if return_times=True) 'estimator': List of fold estimators, length n_cv (if return_estimator=True) 'n_samples': Number of successful folds

RAISES	DESCRIPTION
ValueError	If datasets empty, sample counts mismatch, or devices differ.
TypeError	If n_cv invalid or fit_score_fn not callable.
Exception	Re-raised from fit_score_fn if error_score='raise'.

Notes

Each fold creates a fresh copy of estimator to avoid state leakage between folds. Exceptions during fit/score are handled per error_score parameter.

Source code in spectre/validation/dispatch.py

def dispatch_cross_validation(
    estimator: Any | dict[str, Any],
    datasets: dict[str, DataBatch],
    fit_score_fn: Callable,
    n_cv: int | Any = 5,
    shuffle: bool = False,
    random_state: int | None = None,
    device: str | torch.device | None = None,
    n_jobs: int = 1,
    backend: str = "threading",
    verbose: int = 0,
    pre_dispatch: str = "2*n_jobs",
    return_required: list[str] = ["score_test", "score_train"],
    return_estimator: bool = False,
    return_times: bool = True,
    error_score: float | str = "raise",
    dtype: torch.dtype = torch.float32,
) -> dict[str, Any]:
    """
    Low-level framework for parallel K-fold cross-validation with multiple datasets.

    Handles dataset validation, fold generation, parallel execution, and
    result aggregation for custom cross-validation workflows. Users provide
    a callback function that defines fit/score logic for each fold.

    This is the core dispatch function; for single-dataset CV, use `cross_validate()`.

    Parameters
    ----------
    estimator : Any | dict[str, Any]
        Estimator(s) to validate per fold:

        - Single object: Copied and used for all datasets
        - Dict: Maps dataset names to estimators (keys must match `datasets`)

    datasets : dict[str, DataBatch]
        Multi-dataset mapping with required constraints:

        - Keys: Arbitrary names (e.g., `'X'`, `'Y'`, `'metadata'`)
        - Values: `DataBatch` objects with data, weights, target
        - Constraint: All must have identical sample count

        Example: `{'X': DataBatch(X, X_w, None), 'Y': DataBatch(Y, Y_w, y_target)}`

    fit_score_fn : Callable
        Callback with signature:

        `fit_score_fn(estimator, train_data, test_data) -> dict[str, Tensor]`

        - `train_data`: dict mapping dataset names → `DataBatch` (train split)
        - `test_data`: dict mapping dataset names → `DataBatch` (test split)
        - Returns dict with keys matching `return_required`
        - All returned tensors must be scalars

    n_cv : int | Any, default 5
        CV strategy:

        - `int`: Create `KFold(n_splits=n_cv, ...)`
        - Splitter object: sklearn-compatible with `get_n_splits()` and `split()`

    shuffle : bool, default False
        Whether to shuffle before splitting. Only used when `n_cv` is int.

    random_state : int | None, optional
        Random seed for shuffle. Only used when `n_cv` is int.
        Default is None.

    device : str | torch.device | None, optional
        Device for computation. Default uses device of first dataset.
        CUDA device disables parallelization (n_jobs -> 1).

    n_jobs : int, default 1
        Parallel jobs. -1 uses all processors.
        Automatically set to 1 when using CUDA device.

    backend : str, by default "threading"
        Joblib backend ('threading', 'multiprocessing', 'loky', etc.).

    verbose : int, by default 0
        Verbosity level (0=silent, higher=more output).

    pre_dispatch : str, by default "2*n_jobs"
        Pre-dispatch strategy for joblib queue.

    return_required : list[str], by default ["score_test", "score_train"]
        Keys that `fit_score_fn` must return in its output dict.

    return_estimator : bool, by default False
        Whether to include fitted fold estimators in output.

    return_times : bool, default True
        Whether to include fold execution times in output.

    error_score : float | str, default "raise"
        Behavior on fit/score errors:

        - `'raise'`: Re-raise exception
        - `'warn'`: Assign `nan` and issue warning
        - numeric: Assign that value

    dtype : torch.dtype, default torch.float32
        Data type for error score tensors.

    Returns
    -------
    dict[str, Any]
        Dictionary with keys:

        - `'score_test'`: Test scores per fold, shape `(n_cv,)`
        - `'score_train'`: Train scores per fold, shape `(n_cv,)`
          (only if in `return_required`)
        - `'time'`: Fold times in seconds, length `n_cv`
          (if `return_times=True`)
        - `'estimator'`: List of fold estimators, length `n_cv`
          (if `return_estimator=True`)
        - `'n_samples'`: Number of successful folds

    Raises
    ------
    ValueError
        If datasets empty, sample counts mismatch, or devices differ.
    TypeError
        If `n_cv` invalid or `fit_score_fn` not callable.
    Exception
        Re-raised from `fit_score_fn` if `error_score='raise'`.

    Notes
    -----
    Each fold creates a fresh copy of `estimator` to avoid state leakage between
    folds. Exceptions during fit/score are handled per `error_score` parameter.
    """
    # Validate datasets and get device
    validated_datasets, device, n_samples = _validate_datasets(
        datasets, estimator, device
    )
    primary_name = list(validated_datasets.keys())[0]

    # Disable parallel processing for CUDA
    if device.type == "cuda":
        n_jobs = 1

    # Validate fit_score_fn
    _validate_callable(fit_score_fn, "fit_score_fn")

    # Setup CV splitter. For now all datasets use the same splitter
    if isinstance(n_cv, int):
        cv_splitter = KFold(n_splits=n_cv, shuffle=shuffle, random_state=random_state)
    else:
        cv_splitter = n_cv

    try:
        n_splits = cv_splitter.get_n_splits(validated_datasets[primary_name].data)
    except AttributeError:
        raise TypeError(
            f"`n_cv` must be an int or a CV splitter with `get_n_splits()` method, "
            f"got {type(n_cv).__name__}"
        )

    # Setup error handling
    error_value = _setup_error_value(error_score, device, dtype)

    def _fit_and_score(train_idx, test_idx):
        """Execute fit_score_fn on one fold."""
        try:
            # Handle model copying for CUDA tensors
            estimator_cv = _copy_estimator(estimator, device=device)

            # Split all datasets by indices
            train_data = {}
            test_data = {}
            for name, batch in validated_datasets.items():
                train_data[name] = batch[train_idx]
                test_data[name] = batch[test_idx]

            # Time the fit and score
            start_time = time.time()
            scores = fit_score_fn(estimator_cv, train_data, test_data)
            total_time = time.time() - start_time

            # Validate output
            _validate_score_output(scores, return_required, "fit_score_fn")

            # Extract only the required scores
            result = {key: scores[key] for key in return_required}
            result["time"] = total_time
            result["estimator"] = estimator_cv if return_estimator else None

            # Clean up GPU memory
            if device.type == "cuda" and not return_estimator:
                del estimator_cv
                torch.cuda.empty_cache()

            return result

        except Exception as e:
            return _handle_execution_error(
                e, error_value, error_score, "Fold", return_required
            )

    # Run parallel cross-validation
    parallel = _create_parallel_executor(
        backend, n_jobs, n_splits, verbose, pre_dispatch
    )

    # Get splits (try to use target from primary dataset for stratification)
    try:
        primary_batch_validated = validated_datasets[primary_name]
        if primary_batch_validated.has_target():
            splits = cv_splitter.split(
                primary_batch_validated.data, primary_batch_validated.target
            )
        else:
            splits = cv_splitter.split(primary_batch_validated.data)
    except TypeError:
        raise TypeError(
            f"CV splitter {type(cv_splitter).__name__} requires target parameter "
            "for stratification."
        )

    results = parallel(
        joblib.delayed(_fit_and_score)(train_idx, test_idx)
        for train_idx, test_idx in splits
    )

    # Aggregate results
    return _aggregate_results(
        results, return_required, return_times, return_estimator, estimator=None
    )

dispatch_bootstrap(estimator: Any | dict[str, Any], datasets: dict[str, DataBatch], fit_score_fn: Callable, n_bootstrap: int = 1000, sample_size: float | int | None = None, random_state: int | None = None, device: str | torch.device | None = None, n_jobs: int = 1, backend: str = 'threading', verbose: int = 0, pre_dispatch: str = '2*n_jobs', return_estimator: bool = False, return_times: bool = True, error_score: float | str = 'raise', dtype: torch.dtype = torch.float32) -> dict[str, Any]

Low-level framework for parallel bootstrap resampling with multiple datasets.

Performs random sampling with replacement to resample datasets, scores the estimator on each bootstrap sample, and aggregates results. Provides non-parametric confidence intervals and uncertainty estimates.

This is the core dispatch function; for single-dataset bootstrap, use bootstrap().

PARAMETER	DESCRIPTION
estimator	Fitted estimator(s) to evaluate on bootstrap samples: Single object: Used for all datasets Dict: Maps dataset names to estimators (keys must match datasets) TYPE: Any | dict[str, Any]
datasets	Multi-dataset mapping with required constraints: Keys: Arbitrary names (e.g., 'X', 'Y') Values: DataBatch objects with data, weights, target Constraint: All must have identical sample count TYPE: dict[str, DataBatch]
fit_score_fn	Callback with signature: fit_score_fn(estimator, sample_data) -> dict[str, Tensor] sample_data: dict mapping dataset names → DataBatch (resampled) Must return dict with 'score' key Score should be a scalar Tensor TYPE: Callable
n_bootstrap	Number of bootstrap samples to generate. TYPE: int, by default 1000 DEFAULT: 1000
sample_size	Size of each resample. Default is None (full dataset size, standard bootstrap). float in (0, 1]: Fraction of n_samples int > 0: Absolute number of samples None: Original dataset size TYPE: float | int | None DEFAULT: None
random_state	Random seed for reproducible sampling. TYPE: int | None, optional, by default None DEFAULT: None
device	Device for computation. Default uses device of first dataset. CUDA device disables parallelization (n_jobs -> 1). TYPE: str | device | None DEFAULT: None
n_jobs	Parallel jobs. -1 uses all processors. Automatically set to 1 when using CUDA device. TYPE: int, by default 1 DEFAULT: 1
backend	Joblib backend ('threading', 'multiprocessing', 'loky', etc.). TYPE: str, by default "threading" DEFAULT: 'threading'
verbose	Verbosity level (0=silent, higher=more output). TYPE: int, by default 0 DEFAULT: 0
pre_dispatch	Pre-dispatch strategy for joblib queue. TYPE: str, by default "2*n_jobs" DEFAULT: '2*n_jobs'
return_estimator	Whether to include estimator reference in output. (Estimator not copied per sample, unlike CV) TYPE: bool, by default False DEFAULT: False
return_times	Whether to include bootstrap sample execution times in output. TYPE: bool, by default True DEFAULT: True
error_score	Behavior on score errors: 'raise': Re-raise exception 'warn': Assign nan and issue warning numeric: Assign that value TYPE: float | str DEFAULT: "raise"
dtype	Data type for error score tensors. TYPE: torch.dtype, by default torch.float32 DEFAULT: float32

RETURNS	DESCRIPTION
dict[str, Any]	Dictionary with keys: 'score': Bootstrap scores, shape (n_bootstrap,) 'time': Sample times in seconds, length n_bootstrap (if return_times=True) 'estimator': Estimator reference (if return_estimator=True) 'n_samples': Number of successful bootstrap samples

RAISES	DESCRIPTION
ValueError	If datasets empty, sample counts mismatch, or devices differ.
TypeError	If fit_score_fn not callable.
Exception	Re-raised from fit_score_fn if error_score='raise'.

Notes

Bootstrap sampling with replacement enables estimation of sampling distributions without distributional assumptions. Each sample drawn independently with equal probability, allowing repeats. Estimator is not re-fit; score changes only due to data resampling.

Source code in spectre/validation/dispatch.py

def dispatch_bootstrap(
    estimator: Any | dict[str, Any],
    datasets: dict[str, DataBatch],
    fit_score_fn: Callable,
    n_bootstrap: int = 1000,
    sample_size: float | int | None = None,
    random_state: int | None = None,
    device: str | torch.device | None = None,
    n_jobs: int = 1,
    backend: str = "threading",
    verbose: int = 0,
    pre_dispatch: str = "2*n_jobs",
    return_estimator: bool = False,
    return_times: bool = True,
    error_score: float | str = "raise",
    dtype: torch.dtype = torch.float32,
) -> dict[str, Any]:
    """
    Low-level framework for parallel bootstrap resampling with multiple datasets.

    Performs random sampling with replacement to resample datasets, scores
    the estimator on each bootstrap sample, and aggregates results. Provides
    non-parametric confidence intervals and uncertainty estimates.

    This is the core dispatch function; for single-dataset bootstrap, use `bootstrap()`.

    Parameters
    ----------
    estimator : Any | dict[str, Any]
        Fitted estimator(s) to evaluate on bootstrap samples:

        - Single object: Used for all datasets
        - Dict: Maps dataset names to estimators (keys must match `datasets`)

    datasets : dict[str, DataBatch]
        Multi-dataset mapping with required constraints:

        - Keys: Arbitrary names (e.g., `'X'`, `'Y'`)
        - Values: `DataBatch` objects with data, weights, target
        - Constraint: All must have identical sample count

    fit_score_fn : Callable
        Callback with signature:

        `fit_score_fn(estimator, sample_data) -> dict[str, Tensor]`

        - `sample_data`: dict mapping dataset names → `DataBatch` (resampled)
        - Must return dict with `'score'` key
        - Score should be a scalar `Tensor`

    n_bootstrap : int, by default 1000
        Number of bootstrap samples to generate.

    sample_size : float | int | None, optional
        Size of each resample. Default is None (full dataset size, standard bootstrap).

        - `float` in (0, 1]: Fraction of `n_samples`
        - `int` > 0: Absolute number of samples
        - `None`: Original dataset size

    random_state : int | None, optional, by default None
        Random seed for reproducible sampling.

    device : str | torch.device | None, optional
        Device for computation. Default uses device of first dataset.
        CUDA device disables parallelization (n_jobs -> 1).

    n_jobs : int, by default 1
        Parallel jobs. -1 uses all processors.
        Automatically set to 1 when using CUDA device.

    backend : str, by default "threading"
        Joblib backend ('threading', 'multiprocessing', 'loky', etc.).

    verbose : int, by default 0
        Verbosity level (0=silent, higher=more output).

    pre_dispatch : str, by default "2*n_jobs"
        Pre-dispatch strategy for joblib queue.

    return_estimator : bool, by default False
        Whether to include estimator reference in output.
        (Estimator not copied per sample, unlike CV)

    return_times : bool, by default True
        Whether to include bootstrap sample execution times in output.

    error_score : float | str, default "raise"
        Behavior on score errors:

        - `'raise'`: Re-raise exception
        - `'warn'`: Assign `nan` and issue warning
        - numeric: Assign that value

    dtype : torch.dtype, by default torch.float32
        Data type for error score tensors.

    Returns
    -------
    dict[str, Any]
        Dictionary with keys:

        - `'score'`: Bootstrap scores, shape `(n_bootstrap,)`
        - `'time'`: Sample times in seconds, length `n_bootstrap`
          (if `return_times=True`)
        - `'estimator'`: Estimator reference
          (if `return_estimator=True`)
        - `'n_samples'`: Number of successful bootstrap samples

    Raises
    ------
    ValueError
        If datasets empty, sample counts mismatch, or devices differ.
    TypeError
        If `fit_score_fn` not callable.
    Exception
        Re-raised from `fit_score_fn` if `error_score='raise'`.

    Notes
    -----
    Bootstrap sampling with replacement enables estimation of sampling
    distributions without distributional assumptions. Each sample drawn
    independently with equal probability, allowing repeats. Estimator is
    not re-fit; score changes only due to data resampling.
    """
    # Validate datasets and get device
    validated_datasets, device, n_samples = _validate_datasets(
        datasets, estimator, device
    )

    # Disable parallel processing for CUDA
    if device.type == "cuda":
        n_jobs = 1

    # Validate fit_score_fn
    _validate_callable(fit_score_fn, "fit_score_fn")

    # Validate n_bootstrap
    if not isinstance(n_bootstrap, int) or n_bootstrap < 1:
        raise ValueError(f"`n_bootstrap` must be positive integer, got {n_bootstrap}")

    # Determine bootstrap sample size
    if sample_size is None:
        bootstrap_size = n_samples
    elif isinstance(sample_size, float):
        if not (0 < sample_size <= 1):
            raise ValueError(f"`sample_size` must be in (0, 1], got {sample_size}")
        bootstrap_size = int(sample_size * n_samples)
    elif isinstance(sample_size, int):
        if sample_size < 1:
            raise ValueError(f"`sample_size` must be positive, got {sample_size}")
        bootstrap_size = min(sample_size, n_samples)
    else:
        raise TypeError(
            f"`sample_size` must be float or int, got {type(sample_size).__name__}"
        )

    # Setup error handling
    error_value = _setup_error_value(error_score, device, dtype)

    # Generate all bootstrap indices upfront for reproducibility
    if random_state is not None:
        generator = torch.Generator(device=device).manual_seed(random_state)
    else:
        generator = None

    bootstrap_indices_list = [
        torch.randint(
            0, n_samples, (bootstrap_size,), device=device, generator=generator
        )
        for _ in range(n_bootstrap)
    ]

    def _fit_and_score(bootstrap_idx):
        """Score on one bootstrap sample."""
        try:
            # Create bootstrap sample for all datasets
            sample_data = {}
            for name, batch in validated_datasets.items():
                sample_data[name] = batch[bootstrap_idx]

            # Time the scoring
            start_time = time.time()
            scores = fit_score_fn(estimator, sample_data)
            total_time = time.time() - start_time

            # Validate output
            _validate_score_output(scores, ["score"], "fit_score_fn")

            result = {
                "score": scores["score"],
                "time": total_time,
            }

            return result

        except Exception as e:
            return _handle_execution_error(
                e, error_value, error_score, "Bootstrap sample", ["score"]
            )

    # Run parallel bootstrap
    parallel = _create_parallel_executor(
        backend, n_jobs, n_bootstrap, verbose, pre_dispatch
    )

    results = parallel(
        joblib.delayed(_fit_and_score)(bootstrap_idx)
        for bootstrap_idx in bootstrap_indices_list
    )

    # Aggregate results
    return _aggregate_results(
        results, ["score"], return_times, return_estimator, estimator=estimator
    )