fix: resolve mypy type errors in decomposition, utils, and viz

jspace/decomposition.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
+from typing import cast
+
 import numpy as np
 import scipy.optimize
 import torch
 
 
 def _build_dictionary(V: torch.Tensor) -> np.ndarray:
-    return V.detach().cpu().numpy()
+    return cast(np.ndarray, V.detach().cpu().numpy())
 
 
 def _resolve_V(
@@ -111,14 +113,16 @@ def jspace_occupancy(
         new_error = float(np.linalg.norm(new_residual))
         rand_dir = rng.standard_normal(V_np.shape[1]).astype(np.float32)
         rand_dir /= np.linalg.norm(rand_dir) + 1e-9
-        rand_improve = max(
-            0.0,
-            prev_error - np.linalg.norm(residual - (residual @ rand_dir) * rand_dir),
+        rand_improve = float(
+            max(
+                0.0,
+                float(prev_error) - float(np.linalg.norm(residual - (residual @ rand_dir) * rand_dir)),
+            )
         )
-        rel_improve = (prev_error - new_error) / (prev_error + 1e-9)
-        if rel_improve < threshold * (rand_improve / (prev_error + 1e-9)):
+        rel_improve = float(np.subtract(prev_error, new_error)) / float(prev_error + 1e-9)
+        if rel_improve < threshold * (rand_improve / float(prev_error + 1e-9)):
             return k
         selected = trial
-        prev_error = new_error
+        prev_error = float(new_error)
         residual = new_residual
     return max_k

jspace/utils.py

@@ -7,6 +7,7 @@
 import json
 import os
 from pathlib import Path
+from typing import cast
 
 import numpy as np
 import torch
@@ -122,7 +123,7 @@ def load_lens_layer(cache_dir: Path, layer_idx: int) -> np.ndarray:
     matrix = data["J"]
     shape = tuple(int(dim) for dim in matrix.shape)
     _validate_shape(shape)
-    return np.array(matrix)
+    return cast(np.ndarray, np.asarray(matrix, dtype=np.float32))
 
 
 def lens_cache_exists(cache_dir: Path, layer_indices: list) -> bool:

jspace/viz.py

@@ -8,8 +8,9 @@
 
 import base64
 import html
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence, Sized
 from pathlib import Path
+from typing import cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -114,7 +115,7 @@ def jl_track(iterable: Iterable, description: str) -> Iterable:
         console=get_console(),
         transient=True,
     )
-    total = len(iterable) if hasattr(iterable, "__len__") else None
+    total = len(iterable) if isinstance(iterable, Sized) else None
     with progress:
         task = progress.add_task(description, total=total)
         for item in iterable:
@@ -159,7 +160,7 @@ def plot_cka_block(
         ax.add_patch(rect)
 
     tick_step = max(1, len(layers) // 8)
-    ticks = np.arange(0, len(layers), tick_step)
+    ticks = [int(i) for i in np.arange(0, len(layers), tick_step)]
     ax.set_xticks(ticks, [str(layers[i]) for i in ticks])
     ax.set_yticks(ticks, [str(layers[i]) for i in ticks])
     ax.set_xlabel("Source layer")
@@ -201,7 +202,7 @@ def plot_cka_block(
 
     if out_path is not None:
         fig.savefig(out_path, dpi=200, facecolor=fig.get_facecolor())
-    return fig
+    return cast(Figure, fig)
 
 
 def plot_layer_metrics(
@@ -243,7 +244,7 @@ def plot_layer_metrics(
 
     if out_path is not None:
         fig.savefig(out_path, dpi=200, facecolor=fig.get_facecolor())
-    return fig
+    return cast(Figure, fig)
 
 
 _CSS = (