poetry -> uv + removed hashing in ein.py

gbmi/utils/ein.py

@@ -111,41 +111,42 @@ def _apply_single_dim(
     size: Optional[int] = None,
     device=None,
 ) -> TensorLike:
+    # TODO: This currently slows down computation by 1-2 OOMs on cpu and gpu resp.
     # no_dim is called if the dim we're 'iterating' over isn't in the returned expression
-    c: Dict[str, TensorLike]
-    with tensor_cache.access() as c:
-        key = lambda_hash((f, collect, no_dim, hash(size)))
-        if key in c:
-            return c[key]
-
-        idx = ConstraintTrackingTensor(torch.tensor(0, device=device))
-        reified = f(idx)  # type: ignore
-        if size is None:
-            constraints = getattr(idx, "_constraints", [])
-            if len(constraints) > 1:
-                # TODO: name the dimension argument with the error
-                raise ValueError(
-                    f"Error: incompatible constraints for dimension ({constraints})"
-                )
-            elif len(constraints) == 0:
-                # TODO: introduce warning if we fail
-                size = None
-            else:
-                size = list(constraints)[0]
-
-        dim = dims(sizes=[size])
-        if size is not None:
-            idx = torch.arange(size).to(reified.device if device is None else device)
-            xs = f(idx[dim])  # type: ignore
-        else:
-            xs = f(dim)
-        if isinstance(xs, DTensor) and hash(dim) in [hash(i) for i in xs.dims]:
-            result = collect(xs, dim)
+    # c: Dict[str, TensorLike]
+    # with tensor_cache.access() as c:
+    #     key = lambda_hash((f, collect, no_dim, hash(size)))
+    #     if key in c:
+    #         return c[key]
+
+    idx = ConstraintTrackingTensor(torch.tensor(0, device=device))
+    reified = f(idx)  # type: ignore
+    if size is None:
+        constraints = getattr(idx, "_constraints", [])
+        if len(constraints) > 1:
+            # TODO: name the dimension argument with the error
+            raise ValueError(
+                f"Error: incompatible constraints for dimension ({constraints})"
+            )
+        elif len(constraints) == 0:
+            # TODO: introduce warning if we fail
+            size = None
         else:
-            result = no_dim(xs, dim)
+            size = list(constraints)[0]
+
+    dim = dims(sizes=[size])
+    if size is not None:
+        idx = torch.arange(size).to(reified.device if device is None else device)
+        xs = f(idx[dim])  # type: ignore
+    else:
+        xs = f(dim)
+    if isinstance(xs, DTensor) and hash(dim) in [hash(i) for i in xs.dims]:
+        result = collect(xs, dim)
+    else:
+        result = no_dim(xs, dim)
 
-        c[key] = result
-        return result
+    # c[key] = result
+    return result
 
 
 def _apply(

pyproject.toml

@@ -1,70 +1,71 @@
-[tool.poetry]
-name = "gbmi"
-version = "0.1.0"
-description = "Guarantees-Based Mechanistic Interpretability Experiments and Tools"
-authors = [
-    "Jason Gross <[email protected]>",
-    "Euan Ong <[email protected]>",
-    "Alex Gibson <[email protected]>",
-    "Soufiane Noubir <[email protected]>",
-    "Som Bagchi <[email protected]>",
-    "Chun Hei Yip <[email protected]>",
-    "Rajashree Agrawal <[email protected]>",
+[dependency-groups]
+dev = [
+    "pre-commit<5.0.0,>=4.0.0",
+    "black[jupyter]<25.0.0,>=24.0.0",
+    "flake8<8.0.0,>=7.0.0",
+    "ipykernel<7.0.0,>=6.29.5",
 ]
-readme = "README.md"
-packages = [{ include = "gbmi" }]
-
-[tool.poetry.dependencies]
-python = "^3.10"
-torch = "^2.1.2"
-transformer-lens = "2.9.1"
-einops = "^0.8.0"
-plotly = "^5.18.0"
-jupyterlab = "^4.0.10"
-matplotlib = "^3.9.0"
-jaxtyping = "^0.2.25"
-scipy = "^1.11.4"
-wandb = "^0.18.0"
-pytest = "^8.0.0"
-ipywidgets = "^8.1.1"
-mypy = "^1.8.0"
-types-tqdm = "^4.66.0.5"
-lightning = "^2.1.3"
-rich = "^13.7.0"
-expecttest = "^0.3.0"
-prettyprinter = "^0.18.0"
-kaleido = "0.2.1"
-imageio = "^2.33.1"
-frozendict = "^2.4.0"
-pillow = "^11.0.0"
-scikit-learn = "^1.4.1-post.0"
-tikzplotly = "^0.1.2"
-simple-parsing = "^0.1.5"
-nltk = "^3.8.1"
-seaborn = "^0.13.2"
-pandas = "^2.2.1"
-cloudpickle = "^3.0.0"
-ipympl = "^0.9.3"
-dill = "^0.3.8"
-tqdm-multiprocess = "^0.0.11"
-cirron = "^0.4"
-eintorch = "^0.2.5"
-numpy = "^2.0.1"
-tikzplotlib-patched = "^0.10.1.post13"
-datasets = "^3.0.0"
-circuitsvis = "^1.43.2"
-plotly-gif = "^0.0.4"
-
-[tool.poetry.group.dev.dependencies]
-pre-commit = "^4.0.0"
-black = { extras = ["jupyter"], version = "^24.0.0" }
-flake8 = "^7.0.0"
-ipykernel = "^6.29.5"
-
+includes = ["gbmi"]
 [build-system]
-requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["pdm-backend"]
+build-backend = "pdm.backend"
 
 [tool.black]
 line-length = 88
 target-version = ['py310']
+
+[project]
+authors = [
+    {name = "Jason Gross", email = "[email protected]"},
+    {name = "Euan Ong", email = "[email protected]"},
+    {name = "Alex Gibson", email = "[email protected]"},
+    {name = "Soufiane Noubir", email = "[email protected]"},
+    {name = "Som Bagchi", email = "[email protected]"},
+    {name = "Chun Hei Yip", email = "[email protected]"},
+    {name = "Rajashree Agrawal", email = "[email protected]"},
+]
+requires-python = "<4.0,>=3.10"
+dependencies = [
+    "torch<3.0.0,>=2.1.2",
+    "transformer-lens==2.9.1",
+    "einops<1.0.0,>=0.8.0",
+    "plotly<6.0.0,>=5.18.0",
+    "jupyterlab<5.0.0,>=4.0.10",
+    "matplotlib<4.0.0,>=3.9.0",
+    "jaxtyping<1.0.0,>=0.2.25",
+    "scipy<2.0.0,>=1.11.4",
+    "wandb<1.0.0,>=0.18.0",
+    "pytest<9.0.0,>=8.0.0",
+    "ipywidgets<9.0.0,>=8.1.1",
+    "mypy<2.0.0,>=1.8.0",
+    "types-tqdm<5.0.0.0,>=4.66.0.5",
+    "lightning<3.0.0,>=2.1.3",
+    "rich<14.0.0,>=13.7.0",
+    "expecttest<1.0.0,>=0.3.0",
+    "prettyprinter<1.0.0,>=0.18.0",
+    "kaleido==0.2.1",
+    "imageio<3.0.0,>=2.33.1",
+    "frozendict<3.0.0,>=2.4.0",
+    "pillow<12.0.0,>=11.0.0",
+    "scikit-learn<2.0.0.0,>=1.4.1-post.0",
+    "tikzplotly<1.0.0,>=0.1.2",
+    "simple-parsing<1.0.0,>=0.1.5",
+    "nltk<4.0.0,>=3.8.1",
+    "seaborn<1.0.0,>=0.13.2",
+    "pandas<3.0.0,>=2.2.1",
+    "cloudpickle<4.0.0,>=3.0.0",
+    "ipympl<1.0.0,>=0.9.3",
+    "dill<1.0.0,>=0.3.8",
+    "tqdm-multiprocess<1.0.0,>=0.0.11",
+    "cirron<1.0,>=0.4",
+    "eintorch<1.0.0,>=0.2.5",
+    "numpy<3.0.0,>=2.0.1",
+    "tikzplotlib-patched<1.0.0.0,>=0.10.1.post13",
+    "datasets<4.0.0,>=3.0.0",
+    "circuitsvis<2.0.0,>=1.43.2",
+    "plotly-gif<1.0.0,>=0.0.4",
+]
+name = "gbmi"
+version = "0.1.0"
+description = "Guarantees-Based Mechanistic Interpretability Experiments and Tools"
+readme = "README.md"