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; the sweep layer is neutral.

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 fixed, seeded synthetic classification problem, split per seed.

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-07-14 21:29:47,103][HYDRA] Launching 12 jobs locally

[2026-07-14 21:29:47,103][HYDRA] 	#0 : +C=0.01 +seed=0

[2026-07-14 21:29:47,173][HYDRA] 	#1 : +C=0.01 +seed=1

[2026-07-14 21:29:47,237][HYDRA] 	#2 : +C=0.01 +seed=2

[2026-07-14 21:29:47,300][HYDRA] 	#3 : +C=0.1 +seed=0

[2026-07-14 21:29:47,364][HYDRA] 	#4 : +C=0.1 +seed=1

[2026-07-14 21:29:47,425][HYDRA] 	#5 : +C=0.1 +seed=2

[2026-07-14 21:29:47,488][HYDRA] 	#6 : +C=1.0 +seed=0

[2026-07-14 21:29:47,551][HYDRA] 	#7 : +C=1.0 +seed=1

[2026-07-14 21:29:47,612][HYDRA] 	#8 : +C=1.0 +seed=2

[2026-07-14 21:29:47,675][HYDRA] 	#9 : +C=10.0 +seed=0

[2026-07-14 21:29:47,737][HYDRA] 	#10 : +C=10.0 +seed=1

[2026-07-14 21:29:47,800][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-07-15T01:29:47.145374+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]])

• Indexes: (2)
  • C
    PandasIndex

    PandasIndex(Index([0.01, 0.1, 1.0, 10.0], dtype='float64', name='C'))

  • seed
    PandasIndex

    PandasIndex(Index([0, 1, 2], dtype='int64', name='seed'))

• Attributes: (1)

  provenance :

  {"timestamp": "2026-07-15T01:29:47.145374+00:00", "python": "3.10.19", "platform": "macOS-26.5.1-arm64-arm-64bit", "git": {"sha": null, "dirty": null, "branch": null}, "packages": {"mushin-py": "0.5.0", "torch": "2.12.1", "numpy": "2.2.6", "pytorch-lightning": "2.6.5", "hydra-core": "1.3.4", "hydra-zen": "0.16.0"}, "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.