update to work with integer top_values

app.py

@@ -104,7 +104,7 @@ def make_img(
     img: Image.Image,
     patches: Float[Tensor, " n_patches"],
     *,
-    upper: float | None = None,
+    upper: int | None = None,
 ) -> Image.Image:
     # Resize to 256x256 and crop to 224x224
     resize_size_px = (512, 512)
@@ -241,7 +241,7 @@ def load_tensor(path: str | pathlib.Path) -> Tensor:
 
 
 top_img_i = load_tensor(CWD / "data" / "top_img_i.pt")
-top_values = load_tensor(CWD / "data" / "top_values.pt")
+top_values = load_tensor(CWD / "data" / "top_values_uint8.pt")
 sparsity = load_tensor(CWD / "data" / "sparsity.pt")
 
 
@@ -312,7 +312,7 @@ def get_sae_examples(
             upper = top_values[latent].max().item()
 
         latent_images = [
-            make_img(img, patches, upper=upper)
+            make_img(img, patches.to(float), upper=upper)
             for img, patches in img_patch_pairs[:n_sae_examples]
         ]
 
@@ -359,7 +359,7 @@ def get_modified_dist(
 
     values = torch.tensor(
         [
-            unscaled(float(value), top_values[latent].max().item())
+            unscaled(value, top_values[latent].max().item())
             for value, latent in [
                 (value1, latent1),
                 (value2, latent2),
@@ -392,14 +392,14 @@ def get_modified_dist(
 
 
 @beartype.beartype
-def unscaled(x: float, max_obs: float) -> float:
+def unscaled(x: float | int, max_obs: float | int) -> float:
     """Scale from [-20, 20] to [20 * -max_obs, 20 * max_obs]."""
     return map_range(x, (-20.0, 20.0), (-20.0 * max_obs, 20.0 * max_obs))
 
 
 @beartype.beartype
 def map_range(
-    x: float,
+    x: float | int,
     domain: tuple[float | int, float | int],
     range: tuple[float | int, float | int],
 ):
@@ -415,7 +415,7 @@ def add_highlights(
     img: Image.Image,
     patches: Float[np.ndarray, " n_patches"],
     *,
-    upper: float | None = None,
+    upper: int | None = None,
     opacity: float = 0.9,
 ) -> Image.Image:
     if not len(patches):

requirements.txt

@@ -265,7 +265,7 @@ multidict==6.1.0
     #   yarl
 multiprocess==0.70.16
     # via datasets
-narwhals==1.21.0
+narwhals==1.21.1
     # via
     #   altair
     #   marimo
@@ -414,7 +414,7 @@ pygments==2.19.1
     #   marimo
     #   nbconvert
     #   rich
-pymdown-extensions==10.13
+pymdown-extensions==10.14
     # via marimo
 pyparsing==3.2.1
     # via matplotlib
@@ -483,7 +483,7 @@ saev @ git+https://github.com/samuelstevens/saev@c723ff95462736d907b3c1891d3a149
     # via saev-image-classification (pyproject.toml)
 safehttpx==0.1.6
     # via gradio
-safetensors==0.5.0
+safetensors==0.5.1
     # via
     #   open-clip-torch
     #   timm