Framework-agnostic sweeps (scikit-learn, no torch)¶

MultiRunMetricsWorkflow never inspects your model; it sweeps configurations and collects the dict your task returns. So you can train a scikit-learn model inside task (no torch, no Lightning) and still get a labeled xarray.Dataset back, exactly as with a torch model. mushin needs no scikit-learn integration for this.

  Download this notebook

Requirements: this notebook trains a scikit-learn model. mushin has no scikit-learn dependency (that is the point of this notebook), so install it alongside the plotting extra: %pip install scikit-learn "mushin-py[viz]".

In [ ]:

from __future__ import annotations

import tempfile

from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

from mushin import multirun
from mushin.workflows import MultiRunMetricsWorkflow

from __future__ import annotations import tempfile from sklearn.datasets import make_classification from sklearn.linear_model import LogisticRegression from sklearn.model_selection import train_test_split from mushin import multirun from mushin.workflows import MultiRunMetricsWorkflow

A seeded synthetic classification problem. Each seed draws its own sample and its own train/test split, so the spread across seeds covers both.

In [ ]:

def split(seed: int):
    x, y = make_classification(
        n_samples=400, n_features=20, n_informative=5, random_state=seed
    )
    return train_test_split(x, y, test_size=0.25, random_state=seed)

def split(seed: int): x, y = make_classification( n_samples=400, n_features=20, n_informative=5, random_state=seed ) return train_test_split(x, y, test_size=0.25, random_state=seed)

Sweep the inverse-regularization strength C across seeds; the returned dict populates the dataset: no torch, no checkpoint to save.

In [ ]:

class LogRegCSweep(MultiRunMetricsWorkflow):
    @staticmethod
    def task(C: float, seed: int) -> dict:
        x_train, x_test, y_train, y_test = split(seed)
        model = LogisticRegression(C=C, max_iter=1000, random_state=seed)
        model.fit(x_train, y_train)
        return dict(accuracy=float(model.score(x_test, y_test)))


wf = LogRegCSweep()
wf.run(
    C=multirun([0.01, 0.1, 1.0, 10.0]),
    seed=multirun([0, 1, 2]),
    working_dir=tempfile.mkdtemp(),
)
ds = wf.to_xarray()
ds

class LogRegCSweep(MultiRunMetricsWorkflow): @staticmethod def task(C: float, seed: int) -> dict: x_train, x_test, y_train, y_test = split(seed) model = LogisticRegression(C=C, max_iter=1000, random_state=seed) model.fit(x_train, y_train) return dict(accuracy=float(model.score(x_test, y_test))) wf = LogRegCSweep() wf.run( C=multirun([0.01, 0.1, 1.0, 10.0]), seed=multirun([0, 1, 2]), working_dir=tempfile.mkdtemp(), ) ds = wf.to_xarray() ds

[2026-10-15 01:54:33,062][HYDRA] Launching 12 jobs locally

[2026-10-15 01:54:33,062][HYDRA] 	#0 : +C=0.01 +seed=0

[2026-10-15 01:54:33,140][HYDRA] 	#1 : +C=0.01 +seed=1

[2026-10-15 01:54:33,297][HYDRA] 	#2 : +C=0.01 +seed=2

[2026-10-15 01:54:33,559][HYDRA] 	#3 : +C=0.1 +seed=0

[2026-10-15 01:54:33,659][HYDRA] 	#4 : +C=0.1 +seed=1

[2026-10-15 01:54:33,774][HYDRA] 	#5 : +C=0.1 +seed=2

[2026-10-15 01:54:33,896][HYDRA] 	#6 : +C=1.0 +seed=0

[2026-10-15 01:54:34,007][HYDRA] 	#7 : +C=1.0 +seed=1

[2026-10-15 01:54:34,120][HYDRA] 	#8 : +C=1.0 +seed=2

[2026-10-15 01:54:34,226][HYDRA] 	#9 : +C=10.0 +seed=0

[2026-10-15 01:54:34,496][HYDRA] 	#10 : +C=10.0 +seed=1

[2026-10-15 01:54:34,563][HYDRA] 	#11 : +C=10.0 +seed=2

Out[ ]:

<xarray.Dataset> Size: 152B
Dimensions:   (C: 4, seed: 3)
Coordinates:
  * C         (C) float64 32B 0.01 0.1 1.0 10.0
  * seed      (seed) int64 24B 0 1 2
Data variables:
    accuracy  (C, seed) float64 96B 0.85 0.66 0.92 0.85 ... 0.93 0.85 0.65 0.93
Attributes:
    provenance:  {"timestamp": "2026-10-15T01:54:33.099941+00:00", "python": ...

xarray.Dataset

• Dimensions:
  • C: 4
  • seed: 3
• Coordinates: (2)
  • C
    (C)
    float64
    0.01 0.1 1.0 10.0

    array([ 0.01,  0.1 ,  1.  , 10.  ])

  • seed
    (seed)
    int64
    0 1 2

    array([0, 1, 2])

• Data variables: (1)
  • accuracy
    (C, seed)
    float64
    0.85 0.66 0.92 ... 0.85 0.65 0.93

    array([[0.85, 0.66, 0.92],
           [0.85, 0.67, 0.96],
           [0.85, 0.65, 0.93],
           [0.85, 0.65, 0.93]])

• Attributes: (1)

  provenance :

  {"timestamp": "2026-10-15T01:54:33.099941+00:00", "python": "3.11.17", "platform": "Linux-6.12.67-linuxkit-aarch64-with-glibc2.36", "git": {"sha": "f7c5ae6282122836aed69c12b82d171145ef9c06", "dirty": false, "branch": "HEAD"}, "packages": {"mushin-py": "0.15.0", "torch": "2.13.0", "numpy": "2.4.6", "pytorch-lightning": "2.6.6", "hydra-core": "1.3.7", "hydra-zen": "0.17.0"}, "accelerator": {"cuda": "13.0", "cudnn": 92000, "device": null}, "config": {"C": 0.01, "seed": 0}}

Same labeled dataset as the torch sweep in notebook 01; plot accuracy vs C.

In [ ]:

import matplotlib.pyplot as plt

mean_acc = ds["accuracy"].mean("seed")
fig, ax = plt.subplots(figsize=(5, 3.2))
for s in ds["seed"].values:
    ax.plot(
        ds["C"].values,
        ds["accuracy"].sel(seed=s).values,
        "o-",
        alpha=0.35,
        color="tab:green",
    )
ax.plot(ds["C"].values, mean_acc.values, "ks-", lw=2, label="mean")
ax.set_xscale("log")
ax.set_xlabel("C (inverse regularization strength)")
ax.set_ylabel("held-out accuracy")
ax.set_title("scikit-learn LogisticRegression: accuracy vs C")
ax.legend()
fig.tight_layout()
plt.show()

import matplotlib.pyplot as plt mean_acc = ds["accuracy"].mean("seed") fig, ax = plt.subplots(figsize=(5, 3.2)) for s in ds["seed"].values: ax.plot( ds["C"].values, ds["accuracy"].sel(seed=s).values, "o-", alpha=0.35, color="tab:green", ) ax.plot(ds["C"].values, mean_acc.values, "ks-", lw=2, label="mean") ax.set_xscale("log") ax.set_xlabel("C (inverse regularization strength)") ax.set_ylabel("held-out accuracy") ax.set_title("scikit-learn LogisticRegression: accuracy vs C") ax.legend() fig.tight_layout() plt.show()

See also

• Notebook 01: the identical flow with a torch model.
• Core concepts: why the sweep layer is framework-agnostic.