change build-backend = setuptools -> uv_build

fix type errors

Author: janosh

.pre-commit-config.yaml

@@ -8,14 +8,14 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.12
+    rev: v0.12.12
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-case-conflict
       - id: check-symlinks
@@ -52,5 +52,6 @@ repos:
     hooks:
       - id: ty
         name: ty check
-        entry: ty check .
+        entry: ty check
         language: python
+        additional_dependencies: [ty]

examples/functorch_mlp_ensemble.ipynb

@@ -159,8 +159,6 @@
     "\n",
     "\n",
     "# visualize different noise levels and then train model on the noisiest one\n",
-    "labels: list[Tensor] = []\n",
-    "points: list[Tensor] = []\n",
     "for noise_std in (0, 0.05, 0.1, 0.2):\n",
     "    points, labels = make_spirals(100, noise_std=noise_std)\n",
     "\n",
@@ -223,25 +221,25 @@
     "\n",
     "\n",
     "def train_step_fn(\n",
-    "    weights: Tensor, batch: Tensor, targets: Tensor, lr: float = 0.2\n",
-    ") -> tuple[Tensor, Tensor, Tensor]:\n",
+    "    weights: list[Tensor], batch: Tensor, targets: Tensor, lr: float = 0.2\n",
+    ") -> tuple[Tensor, Tensor, tuple[Tensor, ...]]:\n",
     "    \"\"\"This function performs a single training step.\n",
     "\n",
     "    Args:\n",
-    "        weights (Tensor): Model weights.\n",
+    "        weights (list[Tensor]): Model weights.\n",
     "        batch (Tensor): Mini-batch of training samples.\n",
     "        targets (Tensor): Ground truth labels for the mini-batch.\n",
     "        lr (float, optional): Learning rate. Defaults to 0.2.\n",
     "\n",
     "    Returns:\n",
-    "        tuple[Tensor, Tensor, Tensor]: Loss, accuracy, and updated weights.\n",
+    "        tuple[Tensor, Tensor, tuple[Tensor, ...]]: Loss, accuracy, and updated weights.\n",
     "    \"\"\"\n",
     "\n",
-    "    def compute_loss(weights: Tensor, batch: Tensor, targets: Tensor) -> Tensor:\n",
+    "    def compute_loss(weights: list[Tensor], batch: Tensor, targets: Tensor) -> Tensor:\n",
     "        output = func_model(weights, batch)\n",
     "        return loss_fn(output, targets)\n",
     "\n",
-    "    def accuracy(weights: Tensor, batch: Tensor, targets: Tensor) -> Tensor:\n",
+    "    def accuracy(weights: list[Tensor], batch: Tensor, targets: Tensor) -> Tensor:\n",
     "        output = func_model(weights, batch)\n",
     "        return (output.argmax(dim=1) == targets).float().mean()\n",
     "\n",
@@ -256,7 +254,7 @@
     "\n",
     "    acc = accuracy(new_weights, batch, targets)\n",
     "\n",
-    "    return loss, acc, new_weights"
+    "    return loss, acc, tuple(new_weights)"
    ]
   },
   {
@@ -303,7 +301,7 @@
     "\n",
     "metrics = {}\n",
     "for step in range(n_train_steps):\n",
-    "    loss, acc, weights = train_step_fn(weights, points, labels)\n",
+    "    loss, acc, weights = train_step_fn(list(weights), points, labels)\n",
     "    if step % 100 == 0:\n",
     "        metrics[step] = {\"loss\": loss.item(), \"acc\": acc.item()}\n",
     "\n",
@@ -362,8 +360,9 @@
     "batched_weights = initialize_ensemble(n_models=5)\n",
     "for step in tqdm(range(n_train_steps), desc=\"training MLP ensemble\"):\n",
     "    losses, accuracies, batched_weights = parallel_train_step_fn(\n",
-    "        batched_weights, points, labels\n",
+    "        list(batched_weights), points, labels\n",
     "    )\n",
+    "    batched_weights = list(batched_weights)\n",
     "\n",
     "    loss_dict = {f\"model {idx}\": loss for idx, loss in enumerate(losses, 1)}\n",
     "    writer.add_scalars(\"training/loss\", loss_dict, step)\n",

pyproject.toml

@@ -39,19 +39,20 @@ excel = ["openpyxl"]
 [project.scripts]
 tb-reducer = "tensorboard_reducer:main"
 
-[tool.setuptools.packages]
-find = { include = ["tensorboard_reducer*"], exclude = ["tests*"] }
+[build-system]
+requires = ["uv_build>=0.7.19"]
+build-backend = "uv_build"
 
-[tool.distutils.bdist_wheel]
-universal = true
+[tool.uv.build-backend]
+module-name = "tensorboard_reducer"
+module-root = ""
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "-p no:warnings"
 
 [tool.ruff]
 target-version = "py311"
-include = ["**/pyproject.toml", "*.ipynb", "*.py", "*.pyi"]
 
 [tool.ruff.lint]
 select = ["ALL"]

tensorboard_reducer/event_loader.py

@@ -78,7 +78,7 @@ def reload(self) -> EventAccumulator:
                 self._process_event(event)
         return self
 
-    def _process_event(self, event: Event) -> None:
+    def _process_event(self, event: type[Event]) -> None:
         """Called whenever an event is loaded."""
         if self._first_event_timestamp is None:
             self._first_event_timestamp = event.wall_time

tensorboard_reducer/load.py

@@ -121,7 +121,7 @@ def load_tb_events(
             if handle_dup_steps == "mean":
                 df_scalar = df_scalar.groupby(df_scalar.index).mean()
             elif handle_dup_steps in ("keep-first", "keep-last"):
-                keep = handle_dup_steps.removeprefix(  # ty: ignore[possibly-unbound-attribute]  # noqa: E501
+                keep = handle_dup_steps.removeprefix(  # ty: ignore[possibly-unbound-attribute]
                     "keep-"
                 )
                 df_scalar = df_scalar[~df_scalar.index.duplicated(keep=keep)]

tensorboard_reducer/write.py

@@ -72,10 +72,10 @@ def write_tb_events(
         list[str]: List of paths to the new TensorBoard event files.
     """
     try:
-        from torch.utils.tensorboard import SummaryWriter
+        from torch.utils.tensorboard import SummaryWriter  # noqa: PLC0415
     except ImportError:
         try:
-            from tensorflow.summary import SummaryWriter
+            from tensorflow.summary import SummaryWriter  # noqa: PLC0415
         except ImportError:
             raise ImportError(
                 "Cannot import SummaryWriter from torch nor tensorflow. "
@@ -172,7 +172,8 @@ def write_data_file(
     # names and tag names as 2nd level
     dict_of_dfs = {op: pd.DataFrame(dic) for op, dic in data_to_write.items()}
     df_out = pd.concat(dict_of_dfs, axis=1)
-    df_out.columns = df_out.columns.swaplevel(0, 1)
+    if isinstance(df_out.columns, pd.MultiIndex):
+        df_out.columns = df_out.columns.swaplevel(i=0, j=1)
     df_out.index.name = "step"
 
     # let pandas handle compression inference from extensions (.csv.gz, .json.bz2, etc.)

tests/test_init.py

@@ -2,11 +2,11 @@
 
 import re
 
+import tensorboard_reducer as tbr
+
 
 def test_init_imports() -> None:
     """Test that all expected imports are available from tensorboard_reducer."""
-    import tensorboard_reducer as tbr
-
     assert callable(tbr.load_tb_events)
     assert callable(tbr.reduce_events)
     assert callable(tbr.write_tb_events)
@@ -16,7 +16,4 @@ def test_init_imports() -> None:
 
 def test_init_version() -> None:
     """Test that the version is available when the package is installed."""
-    # Back up the current state
-    import tensorboard_reducer as tbr
-
     assert re.match(r"\d+\.\d+\.\d+", tbr.__version__) is not None

tests/test_main.py

@@ -85,6 +85,7 @@ def test_main_report_version(capsys: pytest.CaptureFixture[str], arg: str) -> No
     with pytest.raises(SystemExit) as exc_info:
         main([arg])
 
+    assert isinstance(exc_info.value, SystemExit)
     assert exc_info.value.code == 0
 
     stdout, stderr = capsys.readouterr()