formatting

Author: samuelstevens

scripts/benchmark.py

@@ -291,8 +291,7 @@ def plot(results: pathlib.Path):
     alt.renderers.enable("png")
 
     df = (
-        pl
-        .read_json(results)
+        pl.read_json(results)
         .select("results")
         .explode("results")
         .unnest("results")
@@ -302,8 +301,7 @@ def plot(results: pathlib.Path):
     title = str(results.parent)
 
     band = (
-        alt
-        .Chart(df)
+        alt.Chart(df)
         .mark_errorband(extent="stdev")  # mean +/- 1 stddev
         .encode(
             alt.X("n_workers", type="quantitative"),
@@ -315,8 +313,7 @@ def plot(results: pathlib.Path):
     )
 
     line = (
-        alt
-        .Chart(df, title=title)
+        alt.Chart(df, title=title)
         .mark_line(point=True)
         .encode(
             alt.X("n_workers", type="quantitative"),

src/saev/data/bird_mae.py

@@ -443,7 +443,6 @@ def _init_weights(self, module: nn.Module) -> None:
     def forward(
         self, input_values: Float[Tensor, "batch 1 512 128"]
     ) -> dict[str, Tensor]:
-
         bsz, c, w, h = input_values.shape
         assert c == 1
         assert w == self.cfg.img_size_x

tests/test_nn_modeling.py

@@ -34,7 +34,12 @@ def sae_cfgs_comprehensive():
     activation_strategy = st.one_of(relu_strategy, topk_strategy, batch_topk_strategy)
 
     return st.builds(
-        lambda d_model, expansion, seed, normalize_w_dec, remove_parallel_grads, activation: (
+        lambda d_model,
+        expansion,
+        seed,
+        normalize_w_dec,
+        remove_parallel_grads,
+        activation: (
             modeling.SparseAutoencoderConfig(
                 d_model=d_model,
                 d_sae=d_model * expansion,