Repository: marimo-team/marimo

Issue:
persisted memoization (ie saving @mo.cache values)
### Description

I'd like to have the control offered by the standard memoization interface of the `mo.cache` decorator, but have the values persist across kernel restarts (as with `marimo.persistent_cache`)

There are a few third-party options that do something similar (at least for pure functions):

- [`joblib.Memory`](https://joblib.readthedocs.io/en/latest/generated/joblib.Memory.html)
- [diskcache's @cache.memoize](https://grantjenks.com/docs/diskcache/tutorial.html#fanoutcache)
- klepto's [inf_cache](https://klepto.readthedocs.io/en/latest/#klepto.inf_cache) with `file_archive`

How does marimo's built-in caching compare to these options?

AFAICT the main difference would be that `mo.cache` takes closed-over dependent variables and/or source code changes into account in the same way `mo.persistent_cache` does, right? Are there any differences between the caching rules of the two marimo interfaces? Or just whether values are persisted? What is the thinking in having the style of cache interface (function memoization vs block-based-with-context-manager) also dictate whether the cache is persisted to disk or not?

### Suggested solution

Possibly just adding a "persisted" flag to `mo.cache`?

Not sure all the details of how caching is implemented internally, but might make sense to re-use one of these external libs in order to not have to fully reinvent the wheel?

### Alternative

Recommend using one of the existing third-party libs

### Additional context

_No response_

Hints from issue comments:
@gabrielgrant: One (potential? hopefully?) difference is that these other solutions seem to have issues with consistently/correctly caching Pandas DataFrames:

https://github.com/joblib/joblib/issues/1611
https://github.com/grantjenks/python-diskcache/issues/314

This is obviously a pretty common use-case in a notebook environment. Is this handled correctly in marimo's cache implementation(s)?@dmadisetti: I've been promising a cache update and here's the branch:

https://github.com/marimo-team/marimo/pull/3480

I think for the pd case, this is something we can explicitly test against, and maybe squeeze into this branch. Some of the hashing mechanism is inspired by job lib. I think marimo should catch this, since it refers directly to the 
I think I'd like to group some of these changes together to prevent repeated cache misses on version updates

Related to this conversation is this discussion:

https://github.com/marimo-team/marimo/discussions/2653

Where we decided `@persistent_cache` should be a drop in for `@cache`@dmadisetti: Just looked at this code again, and the answer is better than joblib, but still not ideal.

marimo has 2 modes of hashing, ContentAddressed (which is the content of the variable) and ExecutionContext (Which is the code run to produce the variable, used for cases where we cannot get a clear hash).

IF a dataframe is all numerical data then the raw, contiguous memory behind the df is used for hashing. If it is not (df contains objects), then the fallback is to ExecutionContext.

More explicit handling could be done, and I don't think column names contribute to the hash in this case.

Solution patch:
diff --git a/docs/_static/js/math.js b/docs/_static/js/math.js
new file mode 100644
index 00000000..822f08a7
--- /dev/null
+++ b/docs/_static/js/math.js
@@ -0,0 +1,53 @@
+(() => {
+  const renderMath = (element) => {
+    const tex = element.textContent || element.innerHTML;
+    if (tex.startsWith("\\(") && tex.endsWith("\\)")) {
+      katex.render(tex.slice(2, -2), element, { displayMode: false });
+    } else if (tex.startsWith("\\[") && tex.endsWith("\\]")) {
+      katex.render(tex.slice(2, -2), element, { displayMode: true });
+    }
+  };
+
+  const renderAllMath = () => {
+    const maths = document.querySelectorAll(
+      ".arithmatex:not([data-processed])",
+    );
+    maths.forEach((element) => {
+      try {
+        renderMath(element);
+        element.setAttribute("data-processed", "true");
+      } catch (error) {
+        console.warn("Failed to render math:", error);
+      }
+    });
+  };
+
+  // Watch for new content
+  const observer = new MutationObserver((mutations) => {
+    const shouldRender = mutations.some((mutation) => {
+      return mutation.addedNodes.length > 0;
+    });
+    if (shouldRender) {
+      renderAllMath();
+    }
+  });
+
+  const init = () => {
+    if (typeof katex === "undefined") {
+      console.warn("KaTeX not loaded");
+      return;
+    }
+
+    renderAllMath();
+    observer.observe(document.body, {
+      childList: true,
+      subtree: true,
+    });
+  };
+
+  if (document.readyState === "loading") {
+    document.addEventListener("DOMContentLoaded", init);
+  } else {
+    init();
+  }
+})();
diff --git a/docs/api/status.md b/docs/api/status.md
index 0186632b..499728e1 100644
--- a/docs/api/status.md
+++ b/docs/api/status.md
@@ -15,7 +15,7 @@ to `tqdm`.
 ```python
 @app.cell
 def __():
-    rerun = mo.ui.button("Rerun")
+    rerun = mo.ui.button(label="Rerun")
     rerun
     return
 
@@ -46,7 +46,7 @@ async def __():
 ```python
 @app.cell
 def __():
-    rerun = mo.ui.button("Rerun")
+    rerun = mo.ui.button(label="Rerun")
     rerun
     return
 
diff --git a/docs/guides/expensive_notebooks.md b/docs/guides/expensive_notebooks.md
index 5290e18f..8d03d135 100644
--- a/docs/guides/expensive_notebooks.md
+++ b/docs/guides/expensive_notebooks.md
@@ -159,7 +159,7 @@ def expensive_component():
     data = db.query("SELECT * FROM data")
     return mo.ui.table(data)
 
-accordion = mo.ui.accordion({
+accordion = mo.accordion({
     "Charts": mo.lazy(expensive_component)
 })
 ```
diff --git a/docs/guides/state.md b/docs/guides/state.md
index 36ca8dea..e01a3c74 100644
--- a/docs/guides/state.md
+++ b/docs/guides/state.md
@@ -62,7 +62,7 @@ But sometimes, you might want interactions to mutate state:
 !!! warning "Use reactive execution for uni-directional flow"
     If you just want the value of a single element to update another element,
     then **you shouldn't use `mo.state`**. Instead, use marimo's built-in
-    reactive execution --- see the [interactivity guide](`../guides/interactivity.md`).
+    reactive execution --- see the [interactivity guide](../guides/interactivity.md).
 
 For cases like these, marimo provides the function [`mo.state()`](../api/state.md),
 which creates a state object and returns a getter and setter function. When you
diff --git a/frontend/package.json b/frontend/package.json
index f751affc..e97e23c6 100644
--- a/frontend/package.json
+++ b/frontend/package.json
@@ -228,7 +228,7 @@
     "tailwindcss": "^3.4.17",
     "turbo": "^2.0.14",
     "typescript": "^5.7.3",
-    "vite": "^5.4.11",
+    "vite": "^5.4.13",
     "vite-tsconfig-paths": "^4.3.2",
     "vitest": "^3.0.2"
   },
diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml
index 2f1c717a..8fba4795 100644
--- a/frontend/pnpm-lock.yaml
+++ b/frontend/pnpm-lock.yaml
@@ -450,7 +450,7 @@ importers:
         version: 8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)
       '@storybook/react-vite':
         specifier: ^8.4.7
-        version: 8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+        version: 8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       '@swc-jotai/react-refresh':
         specifier: ^0.1.1
         version: 0.1.1
@@ -492,10 +492,10 @@ importers:
         version: 7.15.0(eslint@8.57.0)(typescript@5.7.3)
       '@vitejs/plugin-react':
         specifier: ^4.3.4
-        version: 4.3.4(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+        version: 4.3.4(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       '@vitejs/plugin-react-swc':
         specifier: ^3.7.2
-        version: 3.7.2(@swc/helpers@0.5.1)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+        version: 3.7.2(@swc/helpers@0.5.1)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       autoprefixer:
         specifier: ^10.4.20
         version: 10.4.20(postcss@8.5.1)
@@ -581,11 +581,11 @@ importers:
         specifier: ^5.7.3
         version: 5.7.3
       vite:
-        specifier: ^5.4.11
-        version: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+        specifier: ^5.4.13
+        version: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
       vite-tsconfig-paths:
         specifier: ^4.3.2
-        version: 4.3.2(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+        version: 4.3.2(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       vitest:
         specifier: ^3.0.2
         version: 3.0.2(@types/node@20.17.14)(jsdom@24.1.1)(less@4.2.0)(terser@5.37.0)
@@ -8782,8 +8782,8 @@ packages:
       vite:
         optional: true
 
-  vite@5.4.12:
-    resolution: {integrity: sha512-KwUaKB27TvWwDJr1GjjWthLMATbGEbeWYZIbGZ5qFIsgPP3vWzLu4cVooqhm5/Z2SPDUMjyPVjTztm5tYKwQxA==}
+  vite@5.4.14:
+    resolution: {integrity: sha512-EK5cY7Q1D8JNhSaPKVK4pwBFvaTmZxEnoKXLG/U9gmdDcihQGNzFlgIvaxezFR4glP1LsuiedwMBqCXH3wZccA==}
     engines: {node: ^18.0.0 || >=20.0.0}
     hasBin: true
     peerDependencies:
@@ -10241,11 +10241,11 @@ snapshots:
       '@types/istanbul-reports': 1.1.2
       '@types/yargs': 13.0.12
 
-  '@joshwooding/vite-plugin-react-docgen-typescript@0.4.2(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@joshwooding/vite-plugin-react-docgen-typescript@0.4.2(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
       magic-string: 0.27.0
       react-docgen-typescript: 2.2.2(typescript@5.7.3)
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     optionalDependencies:
       typescript: 5.7.3
 
@@ -12446,13 +12446,13 @@ snapshots:
       react: 18.3.1
       react-dom: 18.3.1(react@18.3.1)
 
-  '@storybook/builder-vite@8.4.7(storybook@8.4.7)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@storybook/builder-vite@8.4.7(storybook@8.4.7)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
       '@storybook/csf-plugin': 8.4.7(storybook@8.4.7)
       browser-assert: 1.2.1
       storybook: 8.4.7
       ts-dedent: 2.2.0
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
 
   '@storybook/components@8.4.7(storybook@8.4.7)':
     dependencies:
@@ -12512,11 +12512,11 @@ snapshots:
       react-dom: 18.3.1(react@18.3.1)
       storybook: 8.4.7
 
-  '@storybook/react-vite@8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@storybook/react-vite@8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
-      '@joshwooding/vite-plugin-react-docgen-typescript': 0.4.2(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+      '@joshwooding/vite-plugin-react-docgen-typescript': 0.4.2(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       '@rollup/pluginutils': 5.0.2
-      '@storybook/builder-vite': 8.4.7(storybook@8.4.7)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+      '@storybook/builder-vite': 8.4.7(storybook@8.4.7)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       '@storybook/react': 8.4.7(@storybook/test@8.4.7(storybook@8.4.7))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(storybook@8.4.7)(typescript@5.7.3)
       find-up: 5.0.0
       magic-string: 0.30.17
@@ -12526,7 +12526,7 @@ snapshots:
       resolve: 1.22.8
       storybook: 8.4.7
       tsconfig-paths: 4.2.0
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     transitivePeerDependencies:
       - '@storybook/test'
       - rollup
@@ -13239,21 +13239,21 @@ snapshots:
       '@connectrpc/connect': 1.4.0(@bufbuild/protobuf@1.10.0)
       '@connectrpc/connect-web': 1.4.0(@bufbuild/protobuf@1.10.0)(@connectrpc/connect@1.4.0(@bufbuild/protobuf@1.10.0))
 
-  '@vitejs/plugin-react-swc@3.7.2(@swc/helpers@0.5.1)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@vitejs/plugin-react-swc@3.7.2(@swc/helpers@0.5.1)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
       '@swc/core': 1.7.39(@swc/helpers@0.5.1)
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     transitivePeerDependencies:
       - '@swc/helpers'
 
-  '@vitejs/plugin-react@4.3.4(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@vitejs/plugin-react@4.3.4(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
       '@babel/core': 7.26.0
       '@babel/plugin-transform-react-jsx-self': 7.25.9(@babel/core@7.26.0)
       '@babel/plugin-transform-react-jsx-source': 7.25.9(@babel/core@7.26.0)
       '@types/babel__core': 7.20.5
       react-refresh: 0.14.2
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     transitivePeerDependencies:
       - supports-color
 
@@ -13271,13 +13271,13 @@ snapshots:
       chai: 5.1.2
       tinyrainbow: 2.0.0
 
-  '@vitest/mocker@3.0.2(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
+  '@vitest/mocker@3.0.2(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))':
     dependencies:
       '@vitest/spy': 3.0.2
       estree-walker: 3.0.3
       magic-string: 0.30.17
     optionalDependencies:
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
 
   '@vitest/pretty-format@2.0.5':
     dependencies:
@@ -19359,7 +19359,7 @@ snapshots:
       debug: 4.4.0
       es-module-lexer: 1.6.0
       pathe: 2.0.2
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     transitivePeerDependencies:
       - '@types/node'
       - less
@@ -19371,18 +19371,18 @@ snapshots:
       - supports-color
       - terser
 
-  vite-tsconfig-paths@4.3.2(typescript@5.7.3)(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)):
+  vite-tsconfig-paths@4.3.2(typescript@5.7.3)(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)):
     dependencies:
       debug: 4.4.0
       globrex: 0.1.2
       tsconfck: 3.0.3(typescript@5.7.3)
     optionalDependencies:
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
     transitivePeerDependencies:
       - supports-color
       - typescript
 
-  vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0):
+  vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0):
     dependencies:
       esbuild: 0.21.5
       postcss: 8.5.1
@@ -19396,7 +19396,7 @@ snapshots:
   vitest@3.0.2(@types/node@20.17.14)(jsdom@24.1.1)(less@4.2.0)(terser@5.37.0):
     dependencies:
       '@vitest/expect': 3.0.2
-      '@vitest/mocker': 3.0.2(vite@5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
+      '@vitest/mocker': 3.0.2(vite@5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0))
       '@vitest/pretty-format': 3.0.2
       '@vitest/runner': 3.0.2
       '@vitest/snapshot': 3.0.2
@@ -19412,7 +19412,7 @@ snapshots:
       tinyexec: 0.3.2
       tinypool: 1.0.2
       tinyrainbow: 2.0.0
-      vite: 5.4.12(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
+      vite: 5.4.14(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
       vite-node: 3.0.2(@types/node@20.17.14)(less@4.2.0)(terser@5.37.0)
       why-is-node-running: 2.3.0
     optionalDependencies:
diff --git a/frontend/src/components/find-replace/find-replace.tsx b/frontend/src/components/find-replace/find-replace.tsx
index 63e892ef..9fbd0c7f 100644
--- a/frontend/src/components/find-replace/find-replace.tsx
+++ b/frontend/src/components/find-replace/find-replace.tsx
@@ -56,6 +56,12 @@ export const FindReplace: React.FC = () => {
     return openFindReplacePanel();
   });
 
+  const resetMatches = () => {
+    const matches = getMatches();
+    // False count means an invalid regex
+    setMatches(matches === false ? undefined : matches);
+  };
+
   useEffect(() => {
     if (state.isOpen && findInputRef.current) {
       findInputRef.current.focus(); // Focus the input
@@ -75,9 +81,7 @@ export const FindReplace: React.FC = () => {
       return;
     }
 
-    const matches = getMatches();
-    // False count means an invalid regex
-    setMatches(matches === false ? undefined : matches);
+    resetMatches();
     setGlobalSearchQuery();
   }, [
     // Re-search when any of these change
@@ -205,8 +209,10 @@ export const FindReplace: React.FC = () => {
                 size="xs"
                 variant="outline"
                 className="h-6 text-xs"
-                onClick={() => replaceNext()}
-                disabled={state.findText === "" || state.replaceText === ""}
+                onClick={() => {
+                  replaceNext() && resetMatches();
+                }}
+                disabled={state.findText === ""}
               >
                 Replace Next
               </Button>
@@ -220,6 +226,7 @@ export const FindReplace: React.FC = () => {
                   if (!undo) {
                     return;
                   }
+                  resetMatches();
 
                   // Show toast with undo button
                   const { dismiss } = toast({
@@ -235,7 +242,7 @@ export const FindReplace: React.FC = () => {
                     ),
                   });
                 }}
-                disabled={state.findText === "" || state.replaceText === ""}
+                disabled={state.findText === ""}
               >
                 Replace All
               </Button>
diff --git a/marimo/_cli/envinfo.py b/marimo/_cli/envinfo.py
index dda87c5b..d471bff8 100644
--- a/marimo/_cli/envinfo.py
+++ b/marimo/_cli/envinfo.py
@@ -5,7 +5,8 @@ import platform
 import sys
 from typing import Any, Union, cast
 
-from marimo import __version__
+from marimo import __version__, _loggers
+from marimo._config.manager import get_default_config_manager
 from marimo._utils.health import (
     get_chrome_version,
     get_node_version,
@@ -13,6 +14,8 @@ from marimo._utils.health import (
     get_required_modules_list,
 )
 
+LOGGER = _loggers.marimo_logger()
+
 
 def is_win11() -> bool:
     """
@@ -26,7 +29,19 @@ def is_win11() -> bool:
     return False
 
 
-def get_system_info() -> dict[str, Union[str, dict[str, str]]]:
+def get_experimental_flags() -> dict[str, bool]:
+    try:
+        config = get_default_config_manager(current_path=None).get_config()
+        if "experimental" in config:
+            return config["experimental"]
+        else:
+            return {}
+    except Exception:
+        LOGGER.error("Failed to get experimental flags")
+        return {}
+
+
+def get_system_info() -> dict[str, Union[str, dict[str, Any]]]:
     os_version = platform.release()
     if platform.system() == "Windows" and is_win11():
         os_version = "11"
@@ -49,10 +64,11 @@ def get_system_info() -> dict[str, Union[str, dict[str, str]]]:
 
     requirements = get_required_modules_list()
     optional_deps = get_optional_modules_list()
-
+    experimental = get_experimental_flags()
     return {
         **info,
         "Binaries": binaries,
         "Dependencies": requirements,
         "Optional Dependencies": optional_deps,
+        "Experimental Flags": experimental,
     }
diff --git a/marimo/_plugins/stateless/lazy.py b/marimo/_plugins/stateless/lazy.py
index 69e9ee0e..b3a2120d 100644
--- a/marimo/_plugins/stateless/lazy.py
+++ b/marimo/_plugins/stateless/lazy.py
@@ -40,7 +40,7 @@ class lazy(UIElement[bool, bool]):
 
         Create a lazy-loaded accordion:
         ```python
-        accordion = mo.ui.accordion({"Charts": mo.lazy(expensive_component)})
+        accordion = mo.accordion({"Charts": mo.lazy(expensive_component)})
         ```
 
         Usage with async functions:
diff --git a/marimo/_save/loaders/__init__.py b/marimo/_save/loaders/__init__.py
index da49d65c..801fe157 100644
--- a/marimo/_save/loaders/__init__.py
+++ b/marimo/_save/loaders/__init__.py
@@ -1,6 +1,12 @@
 # Copyright 2024 Marimo. All rights reserved.
-from marimo._save.loaders.loader import Loader
+from marimo._save.loaders.loader import Loader, LoaderPartial, LoaderType
 from marimo._save.loaders.memory import MemoryLoader
 from marimo._save.loaders.pickle import PickleLoader
 
-__all__ = ["Loader", "MemoryLoader", "PickleLoader"]
+__all__ = [
+    "Loader",
+    "LoaderPartial",
+    "LoaderType",
+    "MemoryLoader",
+    "PickleLoader",
+]
diff --git a/marimo/_save/loaders/loader.py b/marimo/_save/loaders/loader.py
index 03e92ce1..f5094d48 100644
--- a/marimo/_save/loaders/loader.py
+++ b/marimo/_save/loaders/loader.py
@@ -3,8 +3,10 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, TypeVar
 
+from marimo._runtime.context import get_context
+from marimo._runtime.state import State
 from marimo._save.cache import CACHE_PREFIX, Cache, CacheType
 
 if TYPE_CHECKING:
@@ -20,6 +22,60 @@ INCONSISTENT_CACHE_BOILER_PLATE = (
 )
 
 
+LoaderType = TypeVar("LoaderType", bound="Loader")
+
+
+class LoaderPartial:
+    """Cache implementation sometimes requires a deferred construction.
+    Moreover, for a cache persistence, we utilize the state registry to store
+    the loader such that the loader object is not actually reconstructed if it
+    does not need to be.
+    """
+
+    def __init__(self, loader_type: type[Loader], **kwargs: Any) -> None:
+        self.loader_type = loader_type
+        self.kwargs = kwargs
+
+    def __call__(self, name: str) -> Loader:
+        try:
+            return self.loader_type(name, **self.kwargs)
+        except TypeError as e:
+            raise TypeError(
+                f"Could not create {self.loader_type} from the construction "
+                f"arguments: [{', '.join(self.kwargs.keys())}]. Consider "
+                "setting these arguments explicitly with "
+                f"{self.loader_type}.partial(needed_arg=value)."
+            ) from e
+
+    def create_or_reconfigure(
+        self, name: str, context: str = "cache_partial"
+    ) -> State[Loader]:
+        ctx = get_context()
+        if ctx.state_registry is None:
+            return State(self(name), _name=name, _context=context)
+
+        loader_state: State[Loader] | None = ctx.state_registry.lookup(
+            name, context=context
+        )
+        if loader_state is None:
+            loader = self(name)
+            # State creation automatically registers itself.
+            return State(loader, _name=name, _context=context)
+        else:
+            loader = loader_state()
+            if isinstance(loader, self.loader_type):
+                # Manually set the attributes of the old loader.
+                # Overriding attr.setter is useful for
+                # managed behavior.
+                for key, value in self.kwargs.items():
+                    setattr(loader, key, value)
+            else:
+                loader = self(name)
+                # Replace the previous loader with the new construction.
+                loader_state._set_value(loader)
+        return loader_state
+
+
 class Loader(ABC):
     """Loaders are responsible for saving and loading persistent caches.
 
@@ -33,6 +89,7 @@ class Loader(ABC):
 
     def __init__(self, name: str) -> None:
         self.name = name
+        self._hits = 0
 
     def build_path(self, hashed_context: str, cache_type: CacheType) -> Path:
         prefix = CACHE_PREFIX.get(cache_type, "U_")
@@ -60,6 +117,7 @@ class Loader(ABC):
         assert set(defs | stateful_refs) == set(loaded.defs), (
             INCONSISTENT_CACHE_BOILER_PLATE
         )
+        self._hits += 1
         return Cache(
             loaded.defs,
             hashed_context,
@@ -69,6 +127,21 @@ class Loader(ABC):
             loaded.meta,
         )
 
+    @property
+    def hits(self) -> int:
+        return self._hits
+
+    @classmethod
+    def partial(cls, **kwargs: Any) -> LoaderPartial:
+        return LoaderPartial(cls, **kwargs)
+
+    @classmethod
+    def cache(cls, *args: Any, **kwargs: Any) -> Any:
+        """General `mo.cache` api for this loader"""
+        from marimo._save.save import cache
+
+        return cache(*args, loader=cls, **kwargs)  # type: ignore
+
     @abstractmethod
     def cache_hit(self, hashed_context: str, cache_type: CacheType) -> bool:
         """Check if cache has been hit given a result hash.
diff --git a/marimo/_save/loaders/memory.py b/marimo/_save/loaders/memory.py
index 2e576d18..de98b425 100644
--- a/marimo/_save/loaders/memory.py
+++ b/marimo/_save/loaders/memory.py
@@ -38,8 +38,7 @@ class MemoryLoader(Loader):
         if self.is_lru:
             self._cache = OrderedDict()
             self._cache_lock = threading.Lock()
-        self.max_size = max_size
-        self.hits = 0
+        self._max_size = max_size
         if cache is not None:
             self._maybe_lock(lambda: self._cache.update(cache))
 
@@ -58,7 +57,6 @@ class MemoryLoader(Loader):
         assert self.cache_hit(hashed_context, cache_type), (
             INCONSISTENT_CACHE_BOILER_PLATE
         )
-        self.hits += 1
         key = self.build_path(hashed_context, cache_type)
         if self.is_lru:
             assert isinstance(self._cache, OrderedDict)
@@ -86,7 +84,7 @@ class MemoryLoader(Loader):
             if self.is_lru:
                 self._cache = OrderedDict(self._cache.items())
                 self._cache_lock = threading.Lock()
-            self.max_size = max_size
+            self._max_size = max_size
             return
         assert isinstance(self._cache, OrderedDict)
         assert self._cache_lock is not None
@@ -94,8 +92,16 @@ class MemoryLoader(Loader):
             self.is_lru = max_size > 0
             if not self.is_lru:
                 self._cache = dict(self._cache.items())
-                self.max_size = max_size
+                self._max_size = max_size
                 return
             while len(self._cache) > max_size:
                 self._cache.popitem(last=False)
-        self.max_size = max_size
+        self._max_size = max_size
+
+    @property
+    def max_size(self) -> int:
+        return self._max_size
+
+    @max_size.setter
+    def max_size(self, value: int) -> None:
+        self.resize(value)
diff --git a/marimo/_save/save.py b/marimo/_save/save.py
index 3159a4ea..40aab292 100644
--- a/marimo/_save/save.py
+++ b/marimo/_save/save.py
@@ -2,11 +2,17 @@
 from __future__ import annotations
 
 import ast
+import functools
 import inspect
 import io
 import os
 import sys
 import traceback
+from collections import abc
+from pathlib import Path
+
+# NB: maxsize follows functools.cache, but renamed max_size outside of drop-in
+# api.
 from sys import maxsize as MAXINT
 from typing import (
     TYPE_CHECKING,
@@ -15,6 +21,8 @@ from typing import (
     Optional,
     Type,
     Union,
+    cast,
+    overload,
 )
 
 from marimo._messaging.tracebacks import write_traceback
@@ -29,7 +37,13 @@ from marimo._save.hash import (
     cache_attempt_from_hash,
     content_cache_attempt_from_base,
 )
-from marimo._save.loaders import Loader, MemoryLoader, PickleLoader
+from marimo._save.loaders import (
+    Loader,
+    LoaderPartial,
+    LoaderType,
+    MemoryLoader,
+    PickleLoader,
+)
 from marimo._utils.variables import is_mangled_local, unmangle_local
 
 # Many assertions are for typing and should always pass. This message is a
@@ -54,33 +68,33 @@ class SkipWithBlock(Exception):
     """Special exception to get around executing the with block body."""
 
 
-class _cache_base(object):
+class _cache_call(object):
     """Like functools.cache but notebook-aware. See `cache` docstring`"""
 
     graph: DirectedGraph
     cell_id: str
     module: ast.Module
     _args: list[str]
-    _loader: Optional[State[MemoryLoader]] = None
+    _loader: Optional[State[Loader]] = None
+    _loader_partial: LoaderPartial
     name: str
     fn: Optional[Callable[..., Any]]
 
     def __init__(
         self,
-        _fn: Optional[Callable[..., Any]] = None,
+        _fn: Optional[Callable[..., Any]],
+        loader_partial: LoaderPartial,
         *,
-        # -1 means unbounded cache
-        maxsize: int = -1,
         pin_modules: bool = False,
         hash_type: str = DEFAULT_HASH,
         # frame_offset is the number of frames the __init__ call is nested
         # with respect to definition of _fn
         frame_offset: int = 0,
     ) -> None:
-        self.max_size = maxsize
         self.pin_modules = pin_modules
         self.hash_type = hash_type
         self._frame_offset = frame_offset
+        self._loader_partial = loader_partial
         if _fn is None:
             self.fn = None
         else:
@@ -90,7 +104,7 @@ class _cache_base(object):
     def hits(self) -> int:
         if self._loader is None:
             return 0
-        return self._loader().hits
+        return self.loader.hits
 
     def _set_context(self, fn: Callable[..., Any]) -> None:
         assert callable(fn), "the provided function must be callable"
@@ -121,7 +135,7 @@ class _cache_base(object):
         # block.
         cell_id = ctx.cell_id or ctx.execution_context.cell_id or ""
         self.scoped_refs = set([f"{ARG_PREFIX}{k}" for k in self._args])
-        # # As are the "locals" not in globals
+        # As are the "locals" not in globals
         self.scoped_refs |= set(f_locals.keys()) - set(ctx.globals.keys())
         # Defined in the cell, and currently available in scope
         self.scoped_refs |= ctx.graph.cells[cell_id].defs & set(
@@ -155,15 +169,14 @@ class _cache_base(object):
         # lifetime of the cache will be tied to the global variable.
         # We can invalidate that by making an invalid namespace.
         if ctx.globals != f_locals:
-            name = name + "*"
+            name = f"{name}*"
 
-        context = "cache"
-        self._loader = ctx.state_registry.lookup(name, context=context)
-        if self._loader is None:
-            loader = MemoryLoader(name, max_size=self.max_size)
-            self._loader = State(loader, _name=name, _context=context)
-        else:
-            self._loader().resize(self.max_size)
+        self._loader = self._loader_partial.create_or_reconfigure(name)
+
+    @property
+    def loader(self) -> Loader:
+        assert self._loader is not None, UNEXPECTED_FAILURE_BOILERPLATE
+        return self._loader()
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         # Capture the deferred call case
@@ -172,6 +185,9 @@ class _cache_base(object):
                 raise TypeError(
                     "cache() takes at most 1 argument (expecting function)"
                 )
+            # Remove the additional frames from singledispatch, because invoking
+            # the function directly.
+            self._frame_offset -= 4
             self._set_context(args[0])
             return self
 
@@ -184,7 +200,7 @@ class _cache_base(object):
         attempt = content_cache_attempt_from_base(
             self.base_block,
             scope,
-            self._loader(),
+            self.loader,
             scoped_refs=self.scoped_refs,
             required_refs=set([f"{ARG_PREFIX}{k}" for k in self._args]),
             as_fn=True,
@@ -197,175 +213,22 @@ class _cache_base(object):
         # stateful variables may be global
         scope = {k: v for k, v in scope.items() if k in attempt.stateful_refs}
         attempt.update(scope, meta={"return": response})
-        self._loader().save_cache(attempt)
+        self.loader.save_cache(attempt)
         return response
 
 
-def cache(
-    _fn: Optional[Callable[..., Any]] = None,
-    *,
-    pin_modules: bool = False,
-) -> _cache_base:
-    """Cache the value of a function based on args and closed-over variables.
-
-    Decorating a function with `@mo.cache` will cache its value based on
-    the function's arguments, closed-over values, and the notebook code.
-
-    **Usage.**
-
-    ```python
-    import marimo as mo
-
-
-    @mo.cache
-    def fib(n):
-        if n <= 1:
-            return n
-        return fib(n - 1) + fib(n - 2)
-    ```
-
-    `mo.cache` is similar to `functools.cache`, but with three key benefits:
-
-    1. `mo.cache` persists its cache even if the cell defining the
-        cached function is re-run, as long as the code defining the function
-        (excluding comments and formatting) has not changed.
-    2. `mo.cache` keys on closed-over values in addition to function arguments,
-        preventing accumulation of hidden state associated with
-        `functools.cache`.
-    3. `mo.cache` does not require its arguments to be
-        hashable (only pickleable), meaning it can work with lists, sets, NumPy
-        arrays, PyTorch tensors, and more.
-
-    `mo.cache` obtains these benefits at the cost of slightly higher overhead
-    than `functools.cache`, so it is best used for expensive functions.
-
-    Like `functools.cache`, `mo.cache` is thread-safe.
-
-    The cache has an unlimited maximum size. To limit the cache size, use
-    `@mo.lru_cache`. `mo.cache` is slightly faster than `mo.lru_cache`, but in
-    most applications the difference is negligible.
-
-    **Args**:
-
-    - `pin_modules`: if True, the cache will be invalidated if module versions
-      differ.
-    """
-    return _cache_base(
-        _fn,
-        maxsize=-1,
-        pin_modules=pin_modules,
-        frame_offset=1,
-    )
-
-
-def lru_cache(
-    _fn: Optional[Callable[..., Any]] = None,
-    *,
-    maxsize: int = 128,
-    pin_modules: bool = False,
-) -> _cache_base:
-    """Decorator for LRU caching the return value of a function.
-
-    `mo.lru_cache` is a version of `mo.cache` with a bounded cache size. As an
-    LRU (Least Recently Used) cache, only the last used `maxsize` values are
-    retained, with the oldest values being discarded. For more information,
-    see the documentation of `mo.cache`.
-
-    **Usage.**
-
-    ```python
-    import marimo as mo
-
-
-    @mo.lru_cache
-    def factorial(n):
-        return n * factorial(n - 1) if n else 1
-    ```
-
-    **Args**:
-
-    - `maxsize`: the maximum number of entries in the cache; defaults to 128.
-      Setting to -1 disables cache limits.
-    - `pin_modules`: if True, the cache will be invalidated if module versions
-      differ.
-    """
-
-    return _cache_base(
-        _fn,
-        maxsize=maxsize,
-        pin_modules=pin_modules,
-        frame_offset=1,
-    )
-
-
-class persistent_cache(object):
-    """Save variables to disk and restore them thereafter.
-
-    The `mo.persistent_cache` context manager lets you delimit a block of code
-    in which variables will be cached to disk when they are first computed. On
-    subsequent runs of the cell, if marimo determines that this block of code
-    hasn't changed and neither has its ancestors, it will restore the variables
-    from disk instead of re-computing them, skipping execution of the block
-    entirely.
-
-    Restoration happens even across notebook runs, meaning you can use
-    `mo.persistent_cache` to make notebooks start *instantly*, with variables
-    that would otherwise be expensive to compute already materialized in
-    memory.
-
-    **Usage.**
-
-    ```python
-    with persistent_cache(name="my_cache"):
-        variable = expensive_function()  # This will be cached to disk.
-        print("hello, cache")  # this will be skipped on cache hits
-    ```
-
-    In this example, `variable` will be cached the first time the block
-    is executed, and restored on subsequent runs of the block. If cache
-    conditions are hit, the contents of `with` block will be skipped on
-    execution. This means that side-effects such as writing to stdout and
-    stderr will be skipped on cache hits.
-
-    For function-level memoization, use `@mo.cache` or `@mo.lru_cache`.
-
-    Note that `mo.state` and `UIElement` changes will also trigger cache
-    invalidation, and be accordingly updated.
-
-    **Warning.** Since context abuses sys frame trace, this may conflict with
-    debugging tools or libraries that also use `sys.settrace`.
-
-    **Args**:
-
-    - `name`: the name of the cache, used to set saving path- to manually
-      invalidate the cache, change the name.
-    - `save_path`: the folder in which to save the cache, defaults to
-      `__marimo__/cache` in the directory of the notebook file
-    - `pin_modules`: if True, the cache will be invalidated if module versions
-      differ between runs, defaults to False.
-    """
-
+class _cache_context(object):
     def __init__(
         self,
         name: str,
+        loader: Loader,
         *,
-        save_path: str | None = None,
         pin_modules: bool = False,
-        _loader: Optional[Loader] = None,
+        hash_type: str = DEFAULT_HASH,
     ) -> None:
         # For an implementation sibling regarding the block skipping, see
         # `withhacks` in pypi.
         self.name = name
-        if save_path is None and (root := notebook_dir()) is not None:
-            save_path = str(root / "__marimo__" / "cache")
-        elif save_path is None:
-            # This can happen if the notebook file is unnamed.
-            save_path = os.path.join("__marimo__", "cache")
-
-        if _loader:
-            self._loader = _loader
-        else:
-            self._loader = PickleLoader(name, save_path)
 
         self._skipped = True
         self._cache: Optional[Cache] = None
@@ -375,6 +238,8 @@ class persistent_cache(object):
         self._body_start: int = MAXINT
         # TODO: Consider having a user level setting.
         self.pin_modules = pin_modules
+        self.hash_type = hash_type
+        self._loader = loader
 
     def __enter__(self) -> Self:
         sys.settrace(lambda *_args, **_keys: None)
@@ -440,9 +305,10 @@ class persistent_cache(object):
                     loader=self._loader,
                     context=pre_module,
                     pin_modules=self.pin_modules,
+                    hash_type=self.hash_type,
                 )
 
-                self.cache_type = self._cache
+                self.cache_type = self._cache.cache_type
                 # Raising on the first valid line, prevents a discrepancy where
                 # whitespace in `With`, changes behavior.
                 self._body_start = save_module.body[0].lineno
@@ -511,3 +377,427 @@ class persistent_cache(object):
             tmpio.seek(0)
             write_traceback(tmpio.read())
         return False
+
+
+# A note on overloading:
+# Single dispatch cues off only the first argument, and expects a similar
+# signature for every overload: https://peps.python.org/pep-0443/
+# However.
+# The context and call APIs are slightly different, so `*` expansions are used
+# to propagate that information down to the actually implementation.
+# As such, we also leverage the `@overload` decorator to provide the correct
+# signature and documentation for the singledispatch entry points.
+
+# mypy also seems to do pretty poorly with this overloading, so there is liberal
+# removal type rules.
+
+
+@functools.singledispatch
+def _cache_invocation(
+    arg: Any,
+    loader: Union[LoaderPartial, Loader, LoaderType],
+    *args: Any,
+    frame_offset: int = 1,
+    **kwargs: Any,
+) -> Union[_cache_call, _cache_context]:
+    del loader, args, kwargs, frame_offset
+    raise TypeError(f"Invalid type for cache: {type(arg)}")
+
+
+def _invoke_call(
+    _fn: Callable[..., Any] | None,
+    loader: Union[LoaderPartial, Loader, LoaderType],
+    *args: Any,
+    frame_offset: int = 1,
+    **kwargs: Any,
+) -> _cache_call:
+    if isinstance(loader, Loader):
+        raise TypeError(
+            "A loader instance cannot be passed to cache directly. "
+            f"Specify a loader type (e.g. `{loader.__class__}`) or a loader "
+            f"partial (e.g. `{loader.__class__}.partial(arg=value)`)."
+        )
+    elif isinstance(loader, type) and issubclass(loader, Loader):
+        cache_args = {
+            "pin_modules": kwargs.pop("pin_modules", False),
+            "hash_type": kwargs.pop("hash_type", DEFAULT_HASH),
+        }
+        loader = cast(Loader, loader).partial(**kwargs)
+        kwargs = cache_args
+
+    if not isinstance(loader, LoaderPartial):
+        raise TypeError(
+            "Invalid loader type. "
+            f"Expected a loader partial, got {type(loader)}."
+        )
+    return _cache_call(
+        _fn, loader, *args, frame_offset=frame_offset + 1, **kwargs
+    )
+
+
+@_cache_invocation.register
+def _invoke_call_none(
+    _fn: None,
+    loader: Union[LoaderPartial, Loader, LoaderType],
+    *args: Any,
+    frame_offset: int = 1,
+    **kwargs: Any,
+) -> _cache_call:
+    return _invoke_call(
+        _fn, loader, *args, frame_offset=frame_offset + 1, **kwargs
+    )
+
+
+@_cache_invocation.register
+def _invoke_call_fn(
+    # mypy would like some generics, but this breaks the singledispatch
+    _fn: abc.Callable,  # type: ignore[type-arg]
+    loader: Union[LoaderPartial, Loader, LoaderType],
+    *args: Any,
+    frame_offset: int = 1,
+    **kwargs: Any,
+) -> _cache_call:
+    return _invoke_call(
+        _fn, loader, *args, frame_offset=frame_offset + 1, **kwargs
+    )
+
+
+@_cache_invocation.register
+def _invoke_context(
+    name: str,
+    loader: Union[LoaderPartial, Loader, LoaderType],
+    *args: Any,
+    frame_offset: int = 1,
+    **kwargs: Any,
+) -> _cache_context:
+    del frame_offset
+
+    if isinstance(loader, LoaderPartial):
+        loader = loader.create_or_reconfigure(name)()
+    elif isinstance(loader, type) and issubclass(loader, Loader):
+        cache_args = {
+            "pin_modules": kwargs.pop("pin_modules", False),
+            "hash_type": kwargs.pop("hash_type", DEFAULT_HASH),
+        }
+        # Create through partial for meaningful error message.
+        loader = (
+            cast(Loader, loader)
+            .partial(**kwargs)
+            .create_or_reconfigure(name)()
+        )
+        kwargs = cache_args
+    return _cache_context(name, loader, *args, **kwargs)
+
+
+@overload
+def cache(
+    fn: Optional[Callable[..., Any]] = None,
+    pin_modules: bool = False,
+    loader: LoaderPartial | LoaderType = MemoryLoader,  # type: ignore[assignment]
+) -> _cache_call: ...
+
+
+@overload
+def cache(
+    name: str,
+    pin_modules: bool = False,
+    loader: LoaderPartial | Loader | LoaderType = MemoryLoader,  # type: ignore[assignment]
+) -> _cache_context: ...
+
+
+def cache(  # type: ignore[misc]
+    name: Union[str, Optional[Callable[..., Any]]] = None,
+    *args: Any,
+    loader: Optional[Union[LoaderPartial, Loader]] = None,
+    _frame_offset: int = 1,
+    _internal_interface_not_for_external_use: None = None,
+    **kwargs: Any,
+) -> Union[_cache_call, _cache_context]:
+    """## Cache the value of a function based on args and closed-over variables.
+
+    Decorating a function with `@mo.cache` will cache its value based on
+    the function's arguments, closed-over values, and the notebook code.
+
+    **Usage.**
+
+    ```python
+    import marimo as mo
+
+
+    @mo.cache
+    def fib(n):
+        if n <= 1:
+            return n
+        return fib(n - 1) + fib(n - 2)
+    ```
+
+    `mo.cache` is similar to `functools.cache`, but with three key benefits:
+
+    1. `mo.cache` persists its cache even if the cell defining the
+        cached function is re-run, as long as the code defining the function
+        (excluding comments and formatting) has not changed.
+    2. `mo.cache` keys on closed-over values in addition to function arguments,
+        preventing accumulation of hidden state associated with
+        `functools.cache`.
+    3. `mo.cache` does not require its arguments to be
+        hashable (only pickleable), meaning it can work with lists, sets, NumPy
+        arrays, PyTorch tensors, and more.
+
+    `mo.cache` obtains these benefits at the cost of slightly higher overhead
+    than `functools.cache`, so it is best used for expensive functions.
+
+    Like `functools.cache`, `mo.cache` is thread-safe.
+
+    The cache has an unlimited maximum size. To limit the cache size, use
+    `@mo.lru_cache`. `mo.cache` is slightly faster than `mo.lru_cache`, but in
+    most applications the difference is negligible.
+
+    Note, `mo.cache` can also be used as a drop in replacement for context block
+    caching like `mo.persistent_cache`.
+
+    **Args**:
+
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ.
+
+    ## Context manager to cache the return value of a block of code.
+
+    The `mo.cache` context manager lets you delimit a block of code in which
+    variables will be cached to memory when they are first computed.
+
+    By default, the cache is stored in memory and is not persisted across kernel
+    runs, for that functionality, refer to `mo.persistent_cache`.
+
+    **Usage.**
+
+    ```python
+    with mo.cache("my_cache") as cache:
+        variable = expensive_function()
+    ```
+
+    **Args**:
+
+    - `name`: the name of the cache, used to set saving path- to manually
+      invalidate the cache, change the name.
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ.
+    - `loader`: the loader to use for the cache, defaults to `MemoryLoader`.
+    """
+    arg = name
+    del name
+
+    if loader is None:
+        loader = MemoryLoader.partial(max_size=-1)
+
+    return _cache_invocation(
+        arg,
+        loader,
+        *args,
+        frame_offset=_frame_offset + 1,
+        **kwargs,
+    )
+
+
+@overload
+def lru_cache(
+    fn: Optional[Callable[..., Any]] = None,
+    maxsize: int = 128,
+    pin_modules: bool = False,
+) -> _cache_call: ...
+
+
+@overload
+def lru_cache(
+    name: str,
+    maxsize: int = 128,
+    pin_modules: bool = False,
+) -> _cache_call: ...
+
+
+def lru_cache(  # type: ignore[misc]
+    name: Union[str, Optional[Callable[..., Any]]] = None,
+    maxsize: int = 128,
+    *args: Any,
+    _internal_interface_not_for_external_use: None = None,
+    **kwargs: Any,
+) -> Union[_cache_call, _cache_context]:
+    """Decorator for LRU caching the return value of a function.
+
+    `mo.lru_cache` is a version of `mo.cache` with a bounded cache size. As an
+    LRU (Least Recently Used) cache, only the last used `maxsize` values are
+    retained, with the oldest values being discarded. For more information,
+    see the documentation of `mo.cache`.
+
+    **Usage.**
+
+    ```python
+    import marimo as mo
+
+
+    @mo.lru_cache
+    def factorial(n):
+        return n * factorial(n - 1) if n else 1
+    ```
+
+    **Args**:
+
+    - `maxsize`: the maximum number of entries in the cache; defaults to 128.
+      Setting to -1 disables cache limits.
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ.
+
+    ## Context manager for LRU caching the return value of a block of code.
+
+    **Args**:
+
+    - `name`: Namespace key for the cache.
+    - `maxsize`: the maximum number of entries in the cache; defaults to 128.
+      Setting to -1 disables cache limits.
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ.
+    """
+    arg = name
+    del name
+
+    if {"loader"} & set(kwargs.keys()):
+        raise ValueError(
+            "loader is not a valid argument "
+            "for lru_cache, use mo.cache instead."
+        )
+
+    return cache(  # type: ignore[no-any-return, call-overload]
+        arg,
+        *args,
+        loader=MemoryLoader.partial(max_size=maxsize),
+        _frame_offset=2,
+        **kwargs,
+    )
+
+
+@overload
+def persistent_cache(
+    name: str,
+    save_path: str | None = None,
+    pin_modules: bool = False,
+) -> _cache_context: ...
+
+
+@overload
+def persistent_cache(
+    fn: Optional[Callable[..., Any]] = None,
+    save_path: str | None = None,
+    pin_modules: bool = False,
+) -> _cache_call: ...
+
+
+def persistent_cache(  # type: ignore[misc]
+    name: Union[str, Optional[Callable[..., Any]]] = None,
+    save_path: str | None = None,
+    *args: Any,
+    _internal_interface_not_for_external_use: None = None,
+    **kwargs: Any,
+) -> Union[_cache_call, _cache_context]:
+    """## Context manager to save variables to disk and restore them thereafter.
+
+    The `mo.persistent_cache` context manager lets you delimit a block of code
+    in which variables will be cached to disk when they are first computed. On
+    subsequent runs of the cell, if marimo determines that this block of code
+    hasn't changed and neither has its ancestors, it will restore the variables
+    from disk instead of re-computing them, skipping execution of the block
+    entirely.
+
+    Restoration happens even across notebook runs, meaning you can use
+    `mo.persistent_cache` to make notebooks start *instantly*, with variables
+    that would otherwise be expensive to compute already materialized in
+    memory.
+
+    **Usage.**
+
+    ```python
+    with persistent_cache(name="my_cache"):
+        variable = expensive_function()  # This will be cached to disk.
+        print("hello, cache")  # this will be skipped on cache hits
+    ```
+
+    In this example, `variable` will be cached the first time the block
+    is executed, and restored on subsequent runs of the block. If cache
+    conditions are hit, the contents of `with` block will be skipped on
+    execution. This means that side-effects such as writing to stdout and
+    stderr will be skipped on cache hits.
+
+    Note that `mo.state` and `UIElement` changes will also trigger cache
+    invalidation, and be accordingly updated.
+
+    **Warning.** Since context abuses sys frame trace, this may conflict with
+    debugging tools or libraries that also use `sys.settrace`.
+
+    **Args**:
+
+    - `name`: the name of the cache, used to set saving path- to manually
+      invalidate the cache, change the name.
+    - `save_path`: the folder in which to save the cache, defaults to
+      `__marimo__/cache` in the directory of the notebook file
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ between runs, defaults to False.
+
+
+    ## Decorator for persistently caching the return value of a function.
+
+    `persistent_cache` can also be used as a drop in function-level memoization
+    for `@mo.cache` or `@mo.lru_cache`. This is much slower than cache, but
+    can be useful for saving function values between kernel restarts. For more
+    details, refer to `mo.cache`.
+
+    **Usage.**
+
+    ```python
+    import marimo as mo
+
+
+    @mo.persistent_cache
+    def my_expensive_function():
+        # Do expensive things
+
+    # or
+
+    @mo.persistent_cache(save_path="my/path/to/cache")
+    def my_expensive_function_cached_in_a_certain_location():
+        # Do expensive things
+    ```
+
+    **Args**:
+
+    - `fn`: the wrapped function if no settings are passed.
+    - `save_path`: the folder in which to save the cache, defaults to
+      `__marimo__/cache` in the directory of the notebook file
+    - `pin_modules`: if True, the cache will be invalidated if module versions
+      differ between runs, defaults to False.
+    """
+
+    arg = name
+    del name
+
+    if {"loader"} & set(kwargs.keys()):
+        raise ValueError(
+            "loader is not a valid argument "
+            "for persistent_cache, use mo.cache instead."
+        )
+
+    if save_path is None and (root := notebook_dir()) is not None:
+        save_path = str(root / "__marimo__" / "cache")
+    elif save_path is None:
+        # This can happen if the notebook file is unnamed.
+        save_path = os.path.join("__marimo__", "cache")
+
+    loader = PickleLoader.partial(save_path=Path(save_path))
+    # Injection hook for testing
+    if "_loader" in kwargs:
+        loader = kwargs.pop("_loader")
+
+    return cache(  # type: ignore[no-any-return, call-overload]
+        arg,
+        *args,
+        loader=loader,
+        _frame_offset=2,
+        **kwargs,
+    )
diff --git a/mkdocs.yml b/mkdocs.yml
index e12641ea..70af9226 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -23,9 +23,12 @@ extra:
     - icon: fontawesome/brands/python
       link: https://pypi.org/project/marimo/
 extra_css:
+  - https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.9.0/katex.min.css
   - stylesheets/extra.css
 extra_javascript:
+  - https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.9.0/katex.min.js
   - _static/js/analytics.js
+  - _static/js/math.js
 markdown_extensions:
   - blocks:MarimoBlocksExtension
   - admonition


Test patch:
diff --git a/tests/_cli/test_envinfo.py b/tests/_cli/test_envinfo.py
index da1816a8..66651ebd 100644
--- a/tests/_cli/test_envinfo.py
+++ b/tests/_cli/test_envinfo.py
@@ -1,4 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
 from marimo._cli.envinfo import (
     get_system_info,
 )
@@ -34,3 +36,4 @@ def test_get_system_info() -> None:
     assert "Python Version" in system_info
     assert "Binaries" in system_info
     assert "Dependencies" in system_info
+    assert "Experimental Flags" in system_info
diff --git a/tests/_save/mocks.py b/tests/_save/mocks.py
index 0b99f337..0ec2c73d 100644
--- a/tests/_save/mocks.py
+++ b/tests/_save/mocks.py
@@ -14,13 +14,13 @@ class MockLoader(Loader):
         data: Optional[Dict[str, Any]] = None,
         stateful_refs: Optional[set[str]] = None,
     ) -> None:
-        self.name = name
         self.save_path = save_path
         self._data = data or {}
         self._cache_hit = data is not None
         self._loaded = False
         self._saved = False
         self._stateful_refs = stateful_refs or set()
+        super().__init__(name)
 
     def cache_hit(self, _hashed_context: str, _cache_type: CacheType) -> bool:
         return self._cache_hit
diff --git a/tests/_save/test_cache.py b/tests/_save/test_cache.py
index 7d1d5d94..beb28806 100644
--- a/tests/_save/test_cache.py
+++ b/tests/_save/test_cache.py
@@ -58,6 +58,26 @@ class TestScriptCache:
 
         app.run()
 
+    @staticmethod
+    def test_cache_loader_api() -> None:
+        app = App()
+        app._anonymous_file = True
+
+        @app.cell
+        def one() -> tuple[int]:
+            from tests._save.mocks import MockLoader
+
+            with MockLoader.cache("one", data={"X": 7, "Y": 8}) as cache:
+                Y = 9
+                X = 10
+            assert X == 7
+            assert cache._cache.defs == {"X": 7, "Y": 8}
+            assert not cache._loader._saved
+            assert cache._loader._loaded
+            return X, Y
+
+        app.run()
+
     @staticmethod
     def test_cache_hit_whitespace() -> None:
         app = App()
@@ -381,6 +401,7 @@ class TestAppCache:
                 ),
             ]
         )
+        assert not k.stderr.messages, k.stderr
         assert k.errors == {}
         assert k.errors == {}
         assert k.globals["X"] == 1
@@ -531,6 +552,62 @@ class TestStateCache:
 
 
 class TestCacheDecorator:
+    async def test_basic_cache_api(
+        self, k: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        await k.run(
+            [
+                exec_req.get(
+                    """
+                    from marimo._save.loaders import MemoryLoader
+
+                    @MemoryLoader.cache
+                    def fib(n):
+                        if n <= 1:
+                            return n
+                        return fib(n - 1) + fib(n - 2)
+
+                    a = fib(5)
+                    b = fib(10)
+                """
+                ),
+            ]
+        )
+
+        assert not k.stderr.messages
+        assert k.globals["fib"].hits == 9
+
+        assert k.globals["a"] == 5
+        assert k.globals["b"] == 55
+
+    async def test_basic_cache_api_with_arg(
+        self, k: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        await k.run(
+            [
+                exec_req.get(
+                    """
+                    from marimo._save.loaders import MemoryLoader
+
+                    @MemoryLoader.cache(max_size=2)
+                    def fib(n):
+                        if n <= 1:
+                            return n
+                        return fib(n - 1) + fib(n - 2)
+
+                    a = fib(5)
+                    b = fib(10)
+                """
+                ),
+            ]
+        )
+
+        assert not k.stderr.messages
+        assert k.globals["fib"].hits == 14
+
+        assert k.globals["a"] == 5
+        assert k.globals["b"] == 55
+
     async def test_basic_cache(
         self, k: Kernel, exec_req: ExecReqProvider
     ) -> None:
@@ -581,7 +658,7 @@ class TestCacheDecorator:
             ]
         )
 
-        assert not len(k.stderr.messages)
+        assert not k.stderr.messages, k.stderr
         # More hits with a smaller cache, because it needs to check the cache
         # more.
         assert k.globals["fib"].hits == 14
@@ -589,6 +666,69 @@ class TestCacheDecorator:
         assert k.globals["a"] == 5
         assert k.globals["b"] == 55
 
+    async def test_lru_cache_default(
+        self, k: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        await k.run(
+            [
+                exec_req.get(
+                    """
+                    from marimo._save.save import lru_cache
+
+                    @lru_cache
+                    def fib(n):
+                        if n <= 1:
+                            return n
+                        return fib(n - 1) + fib(n - 2)
+
+                    a = fib(260)
+                    b = fib(10)
+                """
+                ),
+            ]
+        )
+
+        assert not k.stderr.messages
+        # More hits with a smaller cache, because it needs to check the cache
+        # more. Has 256 entries by default, normal cache hits just 259 times.
+        assert k.globals["fib"].hits == 266
+
+        # A little ridiculous, but still low compute.
+        assert (
+            k.globals["a"]
+            == 971183874599339129547649988289594072811608739584170445
+        )
+        assert k.globals["b"] == 55
+
+    async def test_persistent_cache(
+        self, k: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        await k.run(
+            [
+                exec_req.get(
+                    """
+                    from marimo._save.save import persistent_cache
+                    from marimo._save.loaders import MemoryLoader
+
+                    @persistent_cache(_loader=MemoryLoader)
+                    def fib(n):
+                        if n <= 1:
+                            return n
+                        return fib(n - 1) + fib(n - 2)
+
+                    a = fib(5)
+                    b = fib(10)
+                """
+                ),
+            ]
+        )
+
+        assert not k.stderr.messages
+        assert k.globals["fib"].hits == 9
+
+        assert k.globals["a"] == 5
+        assert k.globals["b"] == 55
+
     async def test_cross_cell_cache(
         self, k: Kernel, exec_req: ExecReqProvider
     ) -> None:
@@ -649,9 +789,9 @@ class TestCacheDecorator:
             [
                 exec_req.get(
                     """
-                             from marimo._save.save import cache
-                             from marimo._runtime.state import state
-                             """
+                    from marimo._save.save import cache
+                    from marimo._runtime.state import state
+                    """
                 ),
                 exec_req.get("""external, setter = state(0)"""),
                 exec_req.get(
@@ -679,7 +819,6 @@ class TestCacheDecorator:
             ]
         )
 
-        assert not k.stdout.messages
         assert not k.stderr.messages
 
         assert k.globals["a"] == 5
@@ -698,10 +837,10 @@ class TestCacheDecorator:
             [
                 exec_req.get(
                     """
-                             from marimo._save.save import cache
-                             from marimo._runtime.state import state
-                             import marimo as mo
-                             """
+                    from marimo._save.save import cache
+                    from marimo._runtime.state import state
+                    import marimo as mo
+                    """
                 ),
                 exec_req.get("slider = mo.ui.slider(0, 1)"),
                 exec_req.get("""external, setter = state("a")"""),
@@ -742,6 +881,49 @@ class TestCacheDecorator:
         # 2 * 9 + 2
         assert k.globals["fib"].hits in (9, 18, 20)
 
+    async def test_rerun_update(
+        self, k: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        await k.run(
+            [
+                exec_req.get("""from marimo._save.save import cache"""),
+                exec_req.get(
+                    """
+                    from marimo._runtime.state import state
+                    """
+                ),
+                exec_req.get("""max_size, setter = state(-1)"""),
+                exec_req.get(
+                    """from marimo._save.loaders import MemoryLoader"""
+                ),
+                exec_req.get("""impure = []"""),
+                exec_req.get("""external = 0;c={}"""),
+                exec_req.get_with_id(
+                    cell_id="0",
+                    code="""
+                    @cache(loader=MemoryLoader.partial(max_size=max_size()))
+                    def fib(n):
+                        if n <= 1:
+                            return n + external
+                        return fib(n - 1) + fib(n - 2)
+                    fib(5)
+                """,
+                ),
+                exec_req.get("""a = fib(5)"""),
+                exec_req.get("""b = fib(10); a; impure.append(b)"""),
+                exec_req.get(
+                    """
+                if len(impure) == 1:
+                    setter(256)
+                    b
+                """
+                ),
+            ]
+        )
+        assert not k.stderr.messages, k.stderr
+        assert not k.stdout.messages, k.stdout
+        assert k.globals["fib"].hits == 10 + 3
+
     async def test_full_scope_utilized(
         self, k: Kernel, exec_req: ExecReqProvider
     ) -> None:
@@ -785,7 +967,6 @@ class TestCacheDecorator:
         )
 
         assert not k.stderr.messages, k.stderr
-        assert not k.stdout.messages, k.stdout
 
         assert len(k.globals["impure"]) == 3
         assert {
@@ -796,16 +977,16 @@ class TestCacheDecorator:
 
         # Same name, but should be under different entries
         assert (
-            k.globals["impure"][0][0]._loader()
-            is not k.globals["impure"][1][0]._loader()
+            k.globals["impure"][0][0].loader
+            is not k.globals["impure"][1][0].loader
         )
 
         assert (
             len(
                 {
-                    *k.globals["impure"][0][0]._loader()._cache.keys(),
-                    *k.globals["impure"][1][0]._loader()._cache.keys(),
-                    *k.globals["impure"][2][0]._loader()._cache.keys(),
+                    *k.globals["impure"][0][0].loader._cache.keys(),
+                    *k.globals["impure"][1][0].loader._cache.keys(),
+                    *k.globals["impure"][2][0].loader._cache.keys(),
                 }
             )
             == 2
@@ -861,7 +1042,6 @@ class TestCacheDecorator:
         )
 
         assert not k.stderr.messages, k.stderr
-        assert not k.stdout.messages, k.stdout
 
         assert len(k.globals["impure"]) == 3
         assert {
@@ -872,16 +1052,16 @@ class TestCacheDecorator:
 
         # Same name, but should be under different entries
         assert (
-            k.globals["impure"][0][0]._loader()
-            is not k.globals["impure"][1][0]._loader()
+            k.globals["impure"][0][0].loader
+            is not k.globals["impure"][1][0].loader
         )
 
         assert (
             len(
                 {
-                    *k.globals["impure"][0][0]._loader()._cache.keys(),
-                    *k.globals["impure"][1][0]._loader()._cache.keys(),
-                    *k.globals["impure"][2][0]._loader()._cache.keys(),
+                    *k.globals["impure"][0][0].loader._cache.keys(),
+                    *k.globals["impure"][1][0].loader._cache.keys(),
+                    *k.globals["impure"][2][0].loader._cache.keys(),
                 }
             )
             == 2