diff --git a/.gitignore b/.gitignore
index d2c130f80a..4c49665315 100644
--- a/.gitignore
+++ b/.gitignore
@@ -65,6 +65,11 @@ coverage.xml
 # gallery notebook downloaded data
 doc/gallery/**/data/
 
+# written on every build by the generate_gallery sphinx extension; a custom
+# thumbnail that overrides extraction is committed with git add -f
+doc/gallery/gallery.rst
+doc/_thumbnails/
+
 # JupyterLab session artifacts
 .jupyter/
 .jupyter_ystore.db
\ No newline at end of file
diff --git a/doc/_thumbnails/rewrites/assumptions.png b/doc/_thumbnails/rewrites/assumptions.png
new file mode 100644
index 0000000000..f7c3583787
Binary files /dev/null and b/doc/_thumbnails/rewrites/assumptions.png differ
diff --git a/doc/gallery/rewrites/assumptions.ipynb b/doc/gallery/rewrites/assumptions.ipynb
new file mode 100644
index 0000000000..43642bf0cc
--- /dev/null
+++ b/doc/gallery/rewrites/assumptions.ipynb
@@ -0,0 +1,1188 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7dbcc6d6",
+   "metadata": {},
+   "source": [
+    "(assumptions)=\n",
+    "\n",
+    "# Structural assumptions and assumption-driven rewrites\n",
+    "\n",
+    ":::{post} Jul 9, 2026\n",
+    ":tags: Graph rewrites, Assumptions, Linear algebra\n",
+    ":category: advanced, explanation\n",
+    ":author: Jesse Grabowski\n",
+    ":::\n",
+    "\n",
+    "PyTensor can track structural properties of symbolic tensors (that a matrix is diagonal, triangular, symmetric, positive-definite, orthogonal, a permutation, and so on) and let graph rewrites use those properties to produce faster compiled functions.\n",
+    "\n",
+    "This is a compiler analysis, not a computer-algebra system. The goal is not to prove theorems, but to let a rewrite replace an expensive operation (a general solve, a dense matmul, a Kronecker product) with a cheaper specialized one, without inserting any runtime checks. Facts are attached to `(variable, property)` pairs inside a {class}`~pytensor.graph.fg.FunctionGraph`, inference is lazy and cached, and an answer of *unknown* is both common and legitimate.\n",
+    "\n",
+    "A fact is declared with a single function, {func}`~pytensor.assumptions.assume`. Propagation, implication, and the rewrites that consume the facts then follow automatically."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "95f81fc1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "import pytensor.tensor as pt\n",
+    "from pytensor import dprint, function\n",
+    "\n",
+    "from pytensor.assumptions import (\n",
+    "    ALL_KEYS,\n",
+    "    DIAGONAL,\n",
+    "    IMPLIES,\n",
+    "    LOWER_TRIANGULAR,\n",
+    "    ORTHOGONAL,\n",
+    "    POSITIVE_DEFINITE,\n",
+    "    SYMMETRIC,\n",
+    "    UPPER_TRIANGULAR,\n",
+    "    AssumptionFeature,\n",
+    "    ConflictingAssumptionsError,\n",
+    "    FactState,\n",
+    "    assume,\n",
+    ")\n",
+    "from pytensor.graph.fg import FunctionGraph\n",
+    "\n",
+    "\n",
+    "def facts(var, *keys):\n",
+    "    \"\"\"Read the facts the system infers about ``var``.\n",
+    "\n",
+    "    Attach an ``AssumptionFeature`` to a throwaway ``FunctionGraph`` and return a\n",
+    "    ``{name: FactState}`` mapping. With no keys, report every registered property.\n",
+    "    \"\"\"\n",
+    "    fg = FunctionGraph(outputs=[var], clone=False)\n",
+    "    feature = AssumptionFeature()\n",
+    "    fg.attach_feature(feature)\n",
+    "    keys = keys or ALL_KEYS\n",
+    "    return {key.name: feature.get(var, key) for key in keys}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f8911749",
+   "metadata": {},
+   "source": [
+    "## Declaring an assumption\n",
+    "\n",
+    "{func}`~pytensor.assumptions.assume` attaches one or more structural assumptions to a tensor. Each keyword is `True` (assert the property holds), `False` (assert it does *not* hold), or omitted (say nothing)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "7d57b55f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "SpecifyAssumptions{diagonal} [id A] a={diag}\n",
+      " └─ x [id B]\n"
+     ]
+    }
+   ],
+   "source": [
+    "x = pt.matrix(\"x\", shape=(3, 3))\n",
+    "x_diag = assume(x, diagonal=True)\n",
+    "\n",
+    "dprint(x_diag, print_assumptions=True);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "ffa58ca0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "f = function([x], x_diag)\n",
+    "value = np.arange(9.0).reshape(3, 3)\n",
+    "\n",
+    "# assume() is a no-op at runtime: the output is exactly the input.\n",
+    "np.allclose(f(value), value)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "94388af8",
+   "metadata": {},
+   "source": [
+    "The output is identical to `x` at runtime. {func}`~pytensor.assumptions.assume` wraps the input in a {class}`~pytensor.assumptions.SpecifyAssumptions` node, a no-op view that carries the declared facts and passes gradients through unchanged. Keeping the assumption in the graph means it survives cloning; a later rewrite drains these markers into the graph's fact cache and removes them.\n",
+    "\n",
+    "{func}`~pytensor.printing.debugprint`, imported here as `dprint`, prints what the system knows about each node as an `a={...}` tag, here `a={diag}`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3cfe1614",
+   "metadata": {},
+   "source": [
+    "## Three-valued logic\n",
+    "\n",
+    "Assumptions use three-valued logic. A fact about a `(variable, property)` pair is one of:\n",
+    "\n",
+    "| `FactState` | meaning |\n",
+    "| --- | --- |\n",
+    "| `TRUE` | the property provably holds |\n",
+    "| `FALSE` | the property provably does *not* hold |\n",
+    "| `UNKNOWN` | the system cannot decide (the default, and very common) |\n",
+    "| `CONFLICT` | contradictory evidence, always an error |\n",
+    "\n",
+    "You rarely construct these yourself. To read the facts the system infers about a variable, attach an {class}`~pytensor.assumptions.AssumptionFeature` to a graph and call `.get` (three-valued) or `.check` (collapses to a plain `bool`, `True` only for `TRUE`). The `facts` helper defined above does exactly this."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "377e647a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'diagonal': <FactState.TRUE: 1>,\n",
+       " 'symmetric': <FactState.TRUE: 1>,\n",
+       " 'positive_definite': <FactState.UNKNOWN: 0>,\n",
+       " 'orthogonal': <FactState.UNKNOWN: 0>}"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "facts(x_diag, DIAGONAL, SYMMETRIC, POSITIVE_DEFINITE, ORTHOGONAL)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0ec6f1cd",
+   "metadata": {},
+   "source": [
+    "The distinction between `UNKNOWN` and `FALSE` matters here. The system cannot determine whether a diagonal matrix is positive-definite (a diagonal matrix may or may not be), so it reports `UNKNOWN`, not `FALSE`. `.check()` collapses both to Python `False`, so use `.get()` when the distinction is relevant.\n",
+    "\n",
+    "We asserted only `diagonal=True`, yet `symmetric` is reported `TRUE`. That follows from the implication system, described next."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d36eae88",
+   "metadata": {},
+   "source": [
+    "## Implications: assert the strongest fact\n",
+    "\n",
+    "Properties are connected by a small, explicit implication lattice. A diagonal matrix is symmetric and both-triangular; a positive-definite matrix is symmetric; a permutation matrix is orthogonal and a selection matrix. The `IMPLIES` registry holds these edges:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "01adcd26",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "diagonal           => lower_triangular, upper_triangular, symmetric\n",
+      "positive_definite  => symmetric\n",
+      "permutation        => selection, orthogonal\n"
+     ]
+    }
+   ],
+   "source": [
+    "for stronger, weaker in IMPLIES.items():\n",
+    "    print(f\"{stronger.name:18s} => {', '.join(w.name for w in weaker)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "bc2cb969",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'diagonal': <FactState.TRUE: 1>,\n",
+       " 'symmetric': <FactState.TRUE: 1>,\n",
+       " 'lower_triangular': <FactState.TRUE: 1>,\n",
+       " 'upper_triangular': <FactState.TRUE: 1>}"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# A single diagonal=True assertion answers all four of these as TRUE.\n",
+    "facts(x_diag, DIAGONAL, SYMMETRIC, LOWER_TRIANGULAR, UPPER_TRIANGULAR)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8fc5c6d2",
+   "metadata": {},
+   "source": [
+    "Implication runs in both directions: forward (a stronger `TRUE` makes the weaker facts `TRUE`) and contrapositive (a weaker `FALSE` makes the stronger fact `FALSE`; a matrix that is not symmetric cannot be diagonal). In practice, assert only the strongest property you know and let the weaker ones follow."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "933eb881",
+   "metadata": {},
+   "source": [
+    "## Conflicts are caught\n",
+    "\n",
+    "Because `assume(..., property=False)` records genuine `FALSE` evidence, asserting something the system can *prove* wrong produces a `CONFLICT`, raised as a {class}`~pytensor.assumptions.ConflictingAssumptionsError` the moment the fact is queried. For example, `pt.eye(5)` is provably diagonal, so asserting that it is *not*:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "cb24e78e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ConflictingAssumptionsError: Conflicting evidence for diagonal on SpecifyAssumptions{!diagonal}.0 from owner-inferred rules.\n"
+     ]
+    }
+   ],
+   "source": [
+    "contradiction = assume(pt.eye(5), diagonal=False)\n",
+    "\n",
+    "try:\n",
+    "    facts(contradiction, DIAGONAL)\n",
+    "except ConflictingAssumptionsError as err:\n",
+    "    print(\"ConflictingAssumptionsError:\", err)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4dddbd6a",
+   "metadata": {},
+   "source": [
+    "## Facts propagate through the graph\n",
+    "\n",
+    "You annotate the inputs, not every intermediate result, and the system carries properties forward. Inference walks the graph inputs-first, and each operation has rules describing what it does to a property: the {func}`~pytensor.tensor.linalg.cholesky` factor of a diagonal matrix is diagonal, a transpose swaps lower- and upper-triangular, and the product of two diagonals is diagonal. Constants are inspected directly, so a literal identity matrix is recognized as diagonal, orthogonal, and a permutation.\n",
+    "\n",
+    "Here the `diagonal` fact flows from `A` through `cholesky` to `L` with no annotation on `L` itself:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "0bf627ff",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Blockwise{Cholesky{lower=True, overwrite_a=False}, (m,m)->(m,m)} [id A] a={diag}\n",
+      " └─ SpecifyAssumptions{diagonal} [id B] a={diag}\n",
+      "    └─ A [id C]\n"
+     ]
+    }
+   ],
+   "source": [
+    "A = assume(pt.matrix(\"A\", shape=(3, 3)), diagonal=True)\n",
+    "L = pt.linalg.cholesky(A)\n",
+    "\n",
+    "dprint(L, print_assumptions=True);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d566e835",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "U    {'lower_triangular': <FactState.UNKNOWN: 0>, 'upper_triangular': <FactState.TRUE: 1>}\n",
+      "U.T  {'lower_triangular': <FactState.TRUE: 1>, 'upper_triangular': <FactState.UNKNOWN: 0>}\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Transposing swaps lower- and upper-triangular.\n",
+    "U = assume(pt.matrix(\"U\", shape=(3, 3)), upper_triangular=True)\n",
+    "\n",
+    "print(f\"{'U':<5}{facts(U, LOWER_TRIANGULAR, UPPER_TRIANGULAR)}\")\n",
+    "print(f\"{'U.T':<5}{facts(U.T, LOWER_TRIANGULAR, UPPER_TRIANGULAR)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bd2a612",
+   "metadata": {},
+   "source": [
+    "## Rewrites that consume assumptions\n",
+    "\n",
+    "The preceding machinery exists to enable rewrites. Compiling a function with `mode=\"FAST_RUN\"` runs rewrites that query the fact cache and specialize the graph; the `SpecifyAssumptions` markers are removed in the process. Each of the four examples below produces a compiled graph in which an expensive operation has been eliminated."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f33e0d68",
+   "metadata": {},
+   "source": [
+    "### Diagonal matmul becomes an elementwise product\n",
+    "\n",
+    "A general matrix multiply is $O(n^3)$. If both operands are diagonal, the product is diagonal and reduces to the elementwise product of the two diagonals. No `Matmul` remains in the compiled graph."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "3df8c988",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "FusedElemwise{Mul} [id A] 3\n",
+      " ├─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id B] 2\n",
+      " │  └─ d1 [id C]\n",
+      " ├─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id D] 1\n",
+      " │  └─ d2 [id E]\n",
+      " ├─ [0 1 2] [id F]\n",
+      " ├─ [0 1 2] [id F]\n",
+      " └─ Alloc [id G] 0\n",
+      "    ├─ 0.0 [id H]\n",
+      "    ├─ 3 [id I]\n",
+      "    └─ 3 [id I]\n",
+      "\n",
+      "Inner graphs:\n",
+      "\n",
+      "FusedElemwise{Mul} [id A]\n",
+      " ← AdvancedSetSubtensor [id J]\n",
+      "    ├─ i4 [id K]\n",
+      "    ├─ Mul [id L]\n",
+      "    │  ├─ i0 [id M]\n",
+      "    │  └─ i1 [id N]\n",
+      "    ├─ i3 [id O]\n",
+      "    └─ i3 [id O]\n"
+     ]
+    }
+   ],
+   "source": [
+    "d1 = pt.matrix(\"d1\", shape=(3, 3))\n",
+    "d2 = pt.matrix(\"d2\", shape=(3, 3))\n",
+    "product = assume(d1, diagonal=True) @ assume(d2, diagonal=True)\n",
+    "\n",
+    "f_product = function([d1, d2], product, mode=\"FAST_RUN\")\n",
+    "dprint(f_product);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3aa5b33c",
+   "metadata": {},
+   "source": [
+    "### Orthogonal $Q Q^\\top \\to I$\n",
+    "\n",
+    "For an orthogonal matrix, $Q Q^\\top$ is the identity. The product reduces to a constant, and the multiply is removed entirely."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "9b8d2584",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCopyOp [id A] 0\n",
+      " └─ [[1. 0. 0. ... 0. 0. 1.]] [id B]\n"
+     ]
+    }
+   ],
+   "source": [
+    "q = pt.matrix(\"q\", shape=(3, 3))\n",
+    "q_orth = assume(q, orthogonal=True)\n",
+    "gram = q_orth @ q_orth.T\n",
+    "\n",
+    "f_gram = function([q], gram, mode=\"FAST_RUN\")\n",
+    "dprint(f_gram);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2d4adfc0",
+   "metadata": {},
+   "source": [
+    "### Positive-definite solve becomes a Cholesky solve\n",
+    "\n",
+    "{func}`~pytensor.tensor.linalg.solve` for a general `A` uses an LU-based solver. If `A` is positive-definite, a Cholesky-based solver is roughly twice as fast and more numerically stable. The assumption selects the specialized path, visible as `assume_a='pos'` on the compiled `Solve`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "aa1be802",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Solve{assume_a='pos', lower=False, b_ndim=1, overwrite_a=False, overwrite_b=False} [id A] 0\n",
+      " ├─ A_pd [id B]\n",
+      " └─ b [id C]\n"
+     ]
+    }
+   ],
+   "source": [
+    "A_pd = pt.matrix(\"A_pd\", shape=(3, 3))\n",
+    "b = pt.vector(\"b\", shape=(3,))\n",
+    "solution = pt.linalg.solve(assume(A_pd, positive_definite=True), b)\n",
+    "\n",
+    "f_solve = function([A_pd, b], solution, mode=\"FAST_RUN\")\n",
+    "dprint(f_solve);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a1fc4cb4",
+   "metadata": {},
+   "source": [
+    "### Kronecker product of diagonals\n",
+    "\n",
+    "The {func}`~pytensor.tensor.linalg.kron` product of two diagonal matrices is itself diagonal, so the dense `KroneckerProduct` op is replaced by a construction from the outer product of the two diagonals. No `KroneckerProduct` remains in the compiled graph."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "093560fd",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "AdvancedSetSubtensor [id A] 7\n",
+      " ├─ Alloc [id B] 6\n",
+      " │  ├─ 0.0 [id C]\n",
+      " │  ├─ 12 [id D]\n",
+      " │  └─ 12 [id D]\n",
+      " ├─ Reshape{1} [id E] 5\n",
+      " │  ├─ Mul [id F] 4\n",
+      " │  │  ├─ ExpandDims{axis=1} [id G] 3\n",
+      " │  │  │  └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id H] 2\n",
+      " │  │  │     └─ k1 [id I]\n",
+      " │  │  └─ ExpandDims{axis=0} [id J] 1\n",
+      " │  │     └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id K] 0\n",
+      " │  │        └─ k2 [id L]\n",
+      " │  └─ [-1] [id M]\n",
+      " ├─ [ 0  1  2 ... 9 10 11] [id N]\n",
+      " └─ [ 0  1  2 ... 9 10 11] [id N]\n"
+     ]
+    }
+   ],
+   "source": [
+    "k1 = pt.matrix(\"k1\", shape=(3, 3))\n",
+    "k2 = pt.matrix(\"k2\", shape=(4, 4))\n",
+    "kron = pt.linalg.kron(assume(k1, diagonal=True), assume(k2, diagonal=True))\n",
+    "\n",
+    "f_kron = function([k1, k2], kron, mode=\"FAST_RUN\")\n",
+    "dprint(f_kron);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e2f1232f",
+   "metadata": {},
+   "source": [
+    "## Worked example: a Gaussian process marginal likelihood\n",
+    "\n",
+    "Gaussian process regression is a natural setting for structural assumptions, because its central object is a covariance matrix that is symmetric and positive-definite by construction.\n",
+    "\n",
+    "For inputs $X$, targets $y$, a kernel $k$, and observation noise $\\sigma^2$, the training covariance is\n",
+    "\n",
+    "$$K = k(X, X) + \\sigma^2 I,$$\n",
+    "\n",
+    "which is symmetric positive-definite. The log marginal likelihood is\n",
+    "\n",
+    "$$\\log p(y \\mid X) = -\\tfrac{1}{2}\\left(y^\\top K^{-1} y + \\log|K| + n \\log 2\\pi\\right),$$\n",
+    "\n",
+    "and a direct translation writes $K^{-1}$ as `inv(K)` and $\\log|K|$ as `slogdet(K)`. This is the pattern GP libraries built on PyTensor use: a kernel annotates every covariance $k(X, X)$ as symmetric and positive-definite when it is constructed, the likelihood and posterior code is written in this naive form, and the assumption system specializes it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "765aeb80",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def linalg_op_counts(fn):\n",
+    "    \"\"\"Count the linear-algebra ops remaining in a compiled function.\"\"\"\n",
+    "    keep = (\"Cholesky\", \"Solve\", \"MatrixInverse\", \"SLogDet\", \"Det\", \"LU\")\n",
+    "    counts = {}\n",
+    "    for node in fn.maker.fgraph.apply_nodes:\n",
+    "        name = type(node.op).__name__\n",
+    "        if any(k in name for k in keep):\n",
+    "            counts[name] = counts.get(name, 0) + 1\n",
+    "    return counts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "7f9ecb20",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def exp_quad_cov(X, ls):\n",
+    "    \"\"\"Squared-exponential (RBF) covariance k(X, X).\"\"\"\n",
+    "    sq_dist = ((X[:, None, :] - X[None, :, :]) ** 2).sum(-1)\n",
+    "    return pt.exp(-0.5 * sq_dist / ls**2)\n",
+    "\n",
+    "\n",
+    "X = pt.matrix(\"X\")\n",
+    "y = pt.vector(\"y\")\n",
+    "ls = pt.scalar(\"ls\")\n",
+    "sigma = pt.scalar(\"sigma\")\n",
+    "n = X.shape[0]\n",
+    "\n",
+    "K = exp_quad_cov(X, ls) + sigma**2 * pt.eye(n)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "66b3d5bb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def marginal_log_likelihood(cov):\n",
+    "    quad = y @ pt.linalg.inv(cov) @ y  # y^T K^-1 y\n",
+    "    _, logdet = pt.linalg.slogdet(cov)  # log|K|\n",
+    "    return -0.5 * (quad + logdet + n * np.log(2 * np.pi))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b9ad04c7",
+   "metadata": {},
+   "source": [
+    "Compiled from the plain covariance `K`, the linear algebra falls back to general routines. PyTensor already avoids forming an explicit inverse, but nothing tells it that `K` is symmetric or positive-definite, so the solve and the log-determinant become two independent general (LU-based) factorizations:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "76f57f51",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'SLogDet': 1, 'Solve': 1}"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "f_plain = function([X, y, ls, sigma], marginal_log_likelihood(K), mode=\"FAST_RUN\")\n",
+    "linalg_op_counts(f_plain)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e455be2c",
+   "metadata": {},
+   "source": [
+    "Now assert what a GP library knows at construction time: the covariance is symmetric and positive-definite. This is a single call, exactly the annotation a kernel attaches to `k(X, X)`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "bd0f9834",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'CholeskySolve': 1, 'Cholesky': 1}"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "K_pd = assume(K, positive_definite=True, symmetric=True)\n",
+    "\n",
+    "f_pd = function([X, y, ls, sigma], marginal_log_likelihood(K_pd), mode=\"FAST_RUN\")\n",
+    "linalg_op_counts(f_pd)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "67cc6e6f",
+   "metadata": {},
+   "source": [
+    "The two general factorizations collapse to a single `Cholesky`. Both the inverse-solve (`CholeskySolve`) and the log-determinant reuse that one factor $L$: since $K = L L^\\top$, the term $K^{-1} y$ is obtained by triangular solves and $\\log|K| = 2\\sum_i \\log L_{ii}$. A Cholesky factorization costs about half an LU factorization and is better conditioned for positive-definite matrices, so the assumption improves both speed and numerical stability."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "09845dc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(array(-13.35336976), array(-13.35336976))"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "rng = np.random.default_rng(0)\n",
+    "X_val = rng.normal(size=(6, 1))\n",
+    "y_val = rng.normal(size=6)\n",
+    "\n",
+    "# Same value, different graph.\n",
+    "f_plain(X_val, y_val, 1.0, 0.5), f_pd(X_val, y_val, 1.0, 0.5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fc504932",
+   "metadata": {},
+   "source": [
+    "### Fitting requires the gradient\n",
+    "\n",
+    "Fitting a GP means optimizing the kernel hyperparameters, which needs the gradient of the marginal likelihood. Taking that gradient and inspecting the compiled graph shows the specialization is not yet complete:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "c63a77c4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'CholeskySolve': 3, 'MatrixInverse': 1, 'Cholesky': 1}"
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "loss = marginal_log_likelihood(K_pd)\n",
+    "grad_ls, grad_sigma = pt.grad(loss, [ls, sigma])\n",
+    "\n",
+    "f_grad = function([X, y, ls, sigma], [loss, grad_ls, grad_sigma], mode=\"FAST_RUN\")\n",
+    "linalg_op_counts(f_grad)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f69cd005",
+   "metadata": {},
+   "source": [
+    "A `MatrixInverse` survives. Differentiating $\\log|K|$ produces a standalone $K^{-1}$, and PyTensor turns an inverse into a solve only when it sits next to a matmul; a bare inverse of a matrix it cannot otherwise tell is positive-definite is left as a general `MatrixInverse`.\n",
+    "\n",
+    "The assumption is still available ({func}`~pytensor.assumptions.check_assumption` returns `True` for that matrix), so a small rewrite closes the gap. When the inverse is applied to a positive-definite matrix, factor it once and solve with {func}`~pytensor.tensor.linalg.cho_solve`:\n",
+    "\n",
+    "$$A^{-1} = (L L^\\top)^{-1}, \\qquad L = \\operatorname{chol}(A).$$\n",
+    "\n",
+    "{func}`~pytensor.tensor.linalg.inv` batches its input, so the op to match is `Blockwise(MatrixInverse)`, spelled `blockwise_of(MatrixInverse)`:\n",
+    "\n",
+    "The next cell is not idempotent: `register_specialize` adds to a global rewrite database, so `f_grad` above is the \"before\" graph only because it was compiled first. Re-running that earlier cell now would compile it with the rewrite applied."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "8b0f7040",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pytensor.graph.rewriting.basic import node_rewriter\n",
+    "from pytensor.tensor.rewriting.basic import register_specialize\n",
+    "from pytensor.tensor.rewriting.blockwise import blockwise_of\n",
+    "from pytensor.tensor.linalg import MatrixInverse, cholesky, cho_solve\n",
+    "from pytensor.assumptions import check_assumption\n",
+    "\n",
+    "\n",
+    "@register_specialize\n",
+    "@node_rewriter([blockwise_of(MatrixInverse)])\n",
+    "def inv_of_psd_to_cho_solve(fgraph, node):\n",
+    "    \"\"\"Replace inv(A) with a Cholesky solve when A is known positive-definite.\"\"\"\n",
+    "    [A] = node.inputs\n",
+    "    if not check_assumption(fgraph, A, POSITIVE_DEFINITE):\n",
+    "        return None\n",
+    "    L = cholesky(A, lower=True)\n",
+    "    identity = pt.eye(A.shape[-1], dtype=A.dtype)\n",
+    "    return [cho_solve((L, True), identity)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "fc91e38b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'CholeskySolve': 4, 'Cholesky': 1}"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "f_grad_fast = function([X, y, ls, sigma], [loss, grad_ls, grad_sigma], mode=\"FAST_RUN\")\n",
+    "\n",
+    "assert np.allclose(\n",
+    "    f_grad(X_val, y_val, 1.0, 0.5),\n",
+    "    f_grad_fast(X_val, y_val, 1.0, 0.5),\n",
+    ")\n",
+    "\n",
+    "linalg_op_counts(f_grad_fast)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2ba0e1b5",
+   "metadata": {},
+   "source": [
+    "The `MatrixInverse` is gone, and the count still shows a single `Cholesky`: the factor introduced by the rewrite merges with the one already built for the forward pass, so the whole loss-and-gradient graph shares one factorization. The gradient values are unchanged.\n",
+    "\n",
+    "This rewrite is exactly what a GP library ships. In `ptgp`, for example, it is registered once as `matrix_inverse_specialize`, alongside companions that lower $\\det(L L^\\top)$ and $\\operatorname{diag}(A A^\\top)$, so a full marginal-likelihood-and-gradient graph compiles to a single Cholesky and no explicit inverse."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dfe3f0fc",
+   "metadata": {},
+   "source": [
+    "## Extending the system\n",
+    "\n",
+    "Constructing an {class}`~pytensor.assumptions.AssumptionKey` registers it. From that point the property is a first-class citizen: {func}`~pytensor.assumptions.assume` accepts it by name, `dprint(..., print_assumptions=True)` reports it, and the rewrite that drains declarations into the fact cache resolves it. Nothing else has to be wired up.\n",
+    "\n",
+    "What remains is to say how the property behaves:\n",
+    "\n",
+    "- {func}`~pytensor.assumptions.register_assumption`: a decorator registering a per-operation inference rule. A rule receives `(key, op, feature, fgraph, node, input_states)` and returns one {class}`~pytensor.assumptions.FactState` per output. Pass `prepend=True` to run ahead of rules already registered for the same pair.\n",
+    "- {func}`~pytensor.assumptions.register_matrix_property_rules`: install the standard rules for a property of the trailing two axes, so a new matrix property survives transposes, reshapes, indexing and broadcasting without writing any of them.\n",
+    "- {func}`~pytensor.assumptions.register_implies`: add edges to the implication lattice.\n",
+    "- {func}`~pytensor.assumptions.register_constant_inference`: infer a fact from the data of a literal constant.\n",
+    "\n",
+    "The example below defines an `INVERTIBLE` property, registers that any identity matrix (`Eye`) is invertible, and records that positive-definiteness implies invertibility."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "3838b4f3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "eye       {'invertible': <FactState.TRUE: 1>}\n",
+      "pos_def   {'invertible': <FactState.TRUE: 1>}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pytensor.assumptions import AssumptionKey, register_assumption, register_implies\n",
+    "from pytensor.tensor.basic import Eye\n",
+    "\n",
+    "INVERTIBLE = AssumptionKey(\"invertible\", \"inv\")\n",
+    "\n",
+    "\n",
+    "@register_assumption(INVERTIBLE, Eye)\n",
+    "def _eye_is_invertible(key, op, feature, fgraph, node, input_states):\n",
+    "    return [FactState.TRUE]\n",
+    "\n",
+    "\n",
+    "# A positive-definite matrix is always invertible.\n",
+    "register_implies(POSITIVE_DEFINITE, INVERTIBLE)\n",
+    "\n",
+    "pos_def = assume(pt.matrix(\"m\", shape=(3, 3)), positive_definite=True)\n",
+    "\n",
+    "# ``eye`` is invertible by the rule above; ``pos_def`` follows from the implication.\n",
+    "print(f\"{'eye':<10}{facts(pt.eye(4), INVERTIBLE)}\")\n",
+    "print(f\"{'pos_def':<10}{facts(pos_def, INVERTIBLE)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "57e0ae55",
+   "metadata": {},
+   "source": [
+    "Because the key registered itself, `assume()` takes `invertible=True` alongside the built-in keywords, and the key can declare and query the fact directly:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "aae66c4f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "assume(m, invertible=True)    True\n",
+      "INVERTIBLE.assume(m)          True\n",
+      "m                             False\n",
+      "SpecifyAssumptions{invertible} [id A] a={inv}\n",
+      " └─ m2 [id B]\n"
+     ]
+    }
+   ],
+   "source": [
+    "m = pt.matrix(\"m2\", shape=(3, 3))\n",
+    "\n",
+    "# The first two spell the same declaration; the third is the undeclared control.\n",
+    "spellings = {\n",
+    "    \"assume(m, invertible=True)\": assume(m, invertible=True),\n",
+    "    \"INVERTIBLE.assume(m)\": INVERTIBLE.assume(m),\n",
+    "    \"m\": m,\n",
+    "}\n",
+    "\n",
+    "for label, var in spellings.items():\n",
+    "    print(f\"{label:<30}{INVERTIBLE.holds(var)}\")\n",
+    "\n",
+    "dprint(INVERTIBLE.assume(m), print_assumptions=True);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d3bc3ba3",
+   "metadata": {},
+   "source": [
+    "### A matrix property: stochastic matrices\n",
+    "\n",
+    "The rows of a row-stochastic matrix (a Markov transition matrix) sum to one, and the product of two such matrices is again stochastic.\n",
+    "\n",
+    "`register_matrix_property_rules` supplies the plumbing (the property survives batch indexing, reshapes, and broadcasting), leaving only the rule specific to this property:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "cf5152c9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "P         True\n",
+      "P @ Q     True\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pytensor.assumptions import register_matrix_property_rules\n",
+    "from pytensor.tensor.math import Dot, Sum\n",
+    "\n",
+    "STOCHASTIC = AssumptionKey(\"stochastic\", \"stoch\")\n",
+    "register_matrix_property_rules(STOCHASTIC)\n",
+    "\n",
+    "\n",
+    "@register_assumption(STOCHASTIC, Dot)\n",
+    "def _product_of_stochastic(key, op, feature, fgraph, node, input_states):\n",
+    "    \"\"\"A product of row-stochastic matrices is row-stochastic.\"\"\"\n",
+    "    if all(state is FactState.TRUE for state in input_states):\n",
+    "        return [FactState.TRUE]\n",
+    "    return [FactState.UNKNOWN]\n",
+    "\n",
+    "\n",
+    "P = pt.matrix(\"P\", shape=(3, 3))\n",
+    "Q = pt.matrix(\"Q\", shape=(3, 3))\n",
+    "two_steps = STOCHASTIC.assume(P) @ STOCHASTIC.assume(Q)\n",
+    "\n",
+    "# ``@`` is Blockwise(Dot); the delegate every key gets forwards to the core op.\n",
+    "print(f\"{'P':<10}{STOCHASTIC.holds(STOCHASTIC.assume(P))}\")\n",
+    "print(f\"{'P @ Q':<10}{STOCHASTIC.holds(two_steps)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "27272598",
+   "metadata": {},
+   "source": [
+    "The rewrite reads the fact and replaces the reduction with a constant. Because the fact reached the product, the whole two-step chain collapses: the compiled graph contains neither the sum nor the matmul."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "b9839f9f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Alloc [id A] 0\n",
+      " ├─ 1.0 [id B]\n",
+      " └─ 3 [id C]\n"
+     ]
+    }
+   ],
+   "source": [
+    "@register_specialize\n",
+    "@node_rewriter([Sum])\n",
+    "def _stochastic_rows_sum_to_one(fgraph, node):\n",
+    "    \"\"\"Replace a row-sum with ones when the matrix is known stochastic.\"\"\"\n",
+    "    [mat] = node.inputs\n",
+    "    if mat.type.ndim < 2 or node.op.axis != (mat.type.ndim - 1,):\n",
+    "        return None\n",
+    "    if not check_assumption(fgraph, mat, STOCHASTIC):\n",
+    "        return None\n",
+    "    return [pt.ones(mat.shape[:-1], dtype=node.outputs[0].dtype)]\n",
+    "\n",
+    "\n",
+    "f_rows = function([P, Q], two_steps.sum(axis=-1), mode=\"FAST_RUN\")\n",
+    "dprint(f_rows);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "3c6edb78",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "stream = np.random.default_rng(1)\n",
+    "rows = stream.dirichlet(np.ones(3), size=3)\n",
+    "cols = stream.dirichlet(np.ones(3), size=3)\n",
+    "\n",
+    "# The constant the rewrite inserted is the value the sum would have computed.\n",
+    "np.allclose(f_rows(rows, cols), (rows @ cols).sum(axis=-1))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97f87dd3",
+   "metadata": {},
+   "source": [
+    "### A property that is not about matrices: sorted vectors\n",
+    "\n",
+    "Nothing about the system assumes a property describes a matrix. A vector known to be nondecreasing (a knot vector, bin edges, a grid of quantile levels) makes `sort` redundant, and a contiguous slice of it is still sorted.\n",
+    "\n",
+    "The two rules below are the whole definition."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "dffe0d7e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCopyOp [id A] 1\n",
+      " └─ Subtensor{start:} [id B] 0\n",
+      "    ├─ grid [id C]\n",
+      "    └─ 1 [id D]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pytensor.tensor.sort import SortOp\n",
+    "from pytensor.tensor.subtensor import Subtensor\n",
+    "\n",
+    "SORTED = AssumptionKey(\"sorted\", \"sort\")\n",
+    "\n",
+    "\n",
+    "@register_assumption(SORTED, Subtensor)\n",
+    "def _slice_of_sorted_is_sorted(key, op, feature, fgraph, node, input_states):\n",
+    "    \"\"\"A contiguous slice of a sorted vector is still sorted.\"\"\"\n",
+    "    if input_states[0] is not FactState.TRUE:\n",
+    "        return [FactState.UNKNOWN]\n",
+    "    if all(isinstance(index, slice) for index in op.idx_list):\n",
+    "        return [FactState.TRUE]\n",
+    "    return [FactState.UNKNOWN]\n",
+    "\n",
+    "\n",
+    "@register_specialize\n",
+    "@node_rewriter([SortOp])\n",
+    "def _sort_of_sorted_is_a_noop(fgraph, node):\n",
+    "    \"\"\"Sorting an already-sorted vector returns it unchanged.\"\"\"\n",
+    "    if not check_assumption(fgraph, node.inputs[0], SORTED):\n",
+    "        return None\n",
+    "    return [node.inputs[0]]\n",
+    "\n",
+    "\n",
+    "raw = pt.vector(\"grid\", shape=(5,))\n",
+    "grid = SORTED.assume(raw)\n",
+    "\n",
+    "# The slice carries the fact, so the sort is dropped.\n",
+    "f_sort = function([raw], pt.sort(grid[1:]), mode=\"FAST_RUN\")\n",
+    "dprint(f_sort);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3434a4f4",
+   "metadata": {},
+   "source": [
+    "## Summary\n",
+    "\n",
+    "- Declare structural facts with {func}`~pytensor.assumptions.assume`. The result is a runtime no-op view of `x`.\n",
+    "- Facts are three-valued (`TRUE`, `FALSE`, `UNKNOWN`), stored per graph, inferred lazily, and linked by a small implication lattice, so asserting the strongest property implies the weaker ones.\n",
+    "- Contradictions raise {class}`~pytensor.assumptions.ConflictingAssumptionsError`.\n",
+    "- Facts propagate through the graph automatically; inspect them with `dprint(..., print_assumptions=True)`.\n",
+    "- Compiling with rewrites turns those facts into faster graphs: diagonal matmuls become elementwise products, positive-definite solves become Cholesky solves, orthogonal Gram products fold to the identity, and Kronecker products of diagonals collapse.\n",
+    "- A new property is a custom {class}`~pytensor.assumptions.AssumptionKey`, and constructing one registers it. {func}`~pytensor.assumptions.register_assumption`, {func}`~pytensor.assumptions.register_matrix_property_rules`, {func}`~pytensor.assumptions.register_implies`, and {func}`~pytensor.assumptions.register_constant_inference` say how it behaves, and a `node_rewriter` that calls {func}`~pytensor.assumptions.check_assumption` turns it into a faster graph. A property need not describe a matrix.\n",
+    "\n",
+    "The built-in properties are `diagonal`, `lower_triangular`, `upper_triangular`, `symmetric`, `positive_definite`, `orthogonal`, `selection`, `permutation`, and `unique_indices`.\n",
+    "\n",
+    "The full API is documented in {ref}`libdoc_assumptions`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9b256bc6",
+   "metadata": {},
+   "source": [
+    "## Authors\n",
+    "\n",
+    "- Authored by Jesse Grabowski in August 2026"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4179c367",
+   "metadata": {},
+   "source": [
+    "## Watermark "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "5c81029f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Last updated: Sun, 23 Aug 2026\n",
+      "\n",
+      "Python implementation: CPython\n",
+      "Python version       : 3.14.7\n",
+      "IPython version      : 9.16.1\n",
+      "\n",
+      "pytensor: 3.3.0+16.g1a4cc21de\n",
+      "\n",
+      "numpy   : 2.5.2\n",
+      "pytensor: 3.3.0+16.g1a4cc21de\n",
+      "\n",
+      "Watermark: 2.6.0\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext watermark\n",
+    "%watermark -n -u -v -iv -w -p pytensor"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0e5e052b",
+   "metadata": {},
+   "source": [
+    ":::{include} ../page_footer.md \n",
+    ":::"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/doc/library/assumptions.rst b/doc/library/assumptions.rst
new file mode 100644
index 0000000000..8ebd73cf7c
--- /dev/null
+++ b/doc/library/assumptions.rst
@@ -0,0 +1,74 @@
+.. _libdoc_assumptions:
+
+==============================================================================
+:mod:`assumptions` -- Structural Assumptions and Assumption-Driven Rewrites
+==============================================================================
+
+.. module:: pytensor.assumptions
+   :platform: Unix, Windows
+   :synopsis: Track structural properties of tensors and let rewrites exploit them
+
+The :mod:`pytensor.assumptions` module records structural facts about symbolic
+tensors -- that a matrix is diagonal, triangular, symmetric, positive-definite --
+so that graph rewrites can replace an expensive operation with a cheaper
+specialized one without inserting runtime checks.
+
+Facts are attached to ``(variable, property)`` pairs inside a
+:class:`~pytensor.graph.fg.FunctionGraph`, inference is lazy and cached, and an
+answer of *unknown* is both common and legitimate.
+
+For a worked introduction, see :doc:`the assumptions gallery notebook
+</gallery/rewrites/assumptions>`.
+
+Declaring assumptions
+=====================
+
+.. autofunction:: pytensor.assumptions.assume
+
+.. autoclass:: pytensor.assumptions.SpecifyAssumptions
+
+Inspecting assumptions
+======================
+
+.. autofunction:: pytensor.assumptions.check_assumption
+
+.. autoclass:: pytensor.assumptions.AssumptionFeature
+   :members: get, check
+
+.. autoclass:: pytensor.assumptions.FactState
+
+.. autoclass:: pytensor.assumptions.ConflictingAssumptionsError
+
+.. autofunction:: pytensor.assumptions.summarize_assumptions
+
+.. autofunction:: pytensor.assumptions.assumption_tags
+
+Properties
+==========
+
+Each property is an :class:`AssumptionKey`. The built-in keys are
+``DIAGONAL``, ``LOWER_TRIANGULAR``, ``UPPER_TRIANGULAR``, ``SYMMETRIC``,
+``POSITIVE_DEFINITE``, ``ORTHOGONAL``, ``SELECTION``, ``PERMUTATION``, and
+``UNIQUE_INDICES``. ``MATRIX_KEYS`` holds the eight that describe a matrix;
+``ALL_KEYS`` is a live view of every registered key, including those added by
+downstream libraries.
+
+.. autoclass:: pytensor.assumptions.AssumptionKey
+   :members: assume, holds
+
+Defining a new property
+=======================
+
+Constructing an :class:`AssumptionKey` registers it, after which
+:func:`assume` accepts it by name and ``debugprint(print_assumptions=True)``
+reports it. The functions below say how the new property behaves.
+
+.. autofunction:: pytensor.assumptions.register_assumption
+
+.. autofunction:: pytensor.assumptions.register_matrix_property_rules
+
+.. autofunction:: pytensor.assumptions.register_universal_assumption
+
+.. autofunction:: pytensor.assumptions.register_implies
+
+.. autofunction:: pytensor.assumptions.register_constant_inference
diff --git a/doc/library/index.rst b/doc/library/index.rst
index 6b5dfe29fe..0f73a090d4 100644
--- a/doc/library/index.rst
+++ b/doc/library/index.rst
@@ -15,6 +15,7 @@ Modules
 .. toctree::
    :maxdepth: 1
 
+   assumptions
    compile/index
    config
    d3viz/index
diff --git a/pytensor/assumptions/__init__.py b/pytensor/assumptions/__init__.py
index 8b580139a8..1739b69bfc 100644
--- a/pytensor/assumptions/__init__.py
+++ b/pytensor/assumptions/__init__.py
@@ -14,10 +14,12 @@
 import pytensor.assumptions.subtensor
 import pytensor.assumptions.symmetric
 import pytensor.assumptions.triangular
+from pytensor.assumptions.bundles import register_matrix_property_rules
 from pytensor.assumptions.core import (
     ALL_KEYS,
     DIAGONAL,
     IMPLIES,
+    KEY_REGISTRY,
     LOWER_TRIANGULAR,
     MATRIX_KEYS,
     ORTHOGONAL,
@@ -35,6 +37,7 @@
     register_assumption,
     register_constant_inference,
     register_implies,
+    register_universal_assumption,
 )
 from pytensor.assumptions.specify import (
     SpecifyAssumptions,
diff --git a/pytensor/assumptions/alloc.py b/pytensor/assumptions/alloc.py
index a84c84ce83..35c498aea2 100644
--- a/pytensor/assumptions/alloc.py
+++ b/pytensor/assumptions/alloc.py
@@ -1,5 +1,5 @@
 from pytensor.assumptions.core import (
-    ALL_KEYS,
+    MATRIX_KEYS,
     FactState,
     register_assumption,
     true_if,
@@ -124,5 +124,5 @@ def alloc_propagates_matrix_property(
     return [FactState.UNKNOWN]
 
 
-for _key in ALL_KEYS:
+for _key in MATRIX_KEYS:
     register_assumption(_key, Alloc)(alloc_propagates_matrix_property)
diff --git a/pytensor/assumptions/blockwise.py b/pytensor/assumptions/blockwise.py
index 460cf8e3eb..69cba2cdc0 100644
--- a/pytensor/assumptions/blockwise.py
+++ b/pytensor/assumptions/blockwise.py
@@ -1,17 +1,13 @@
 from pytensor.assumptions.core import (
-    ALL_KEYS,
     infer_assumption_for_node,
-    register_assumption,
+    register_universal_assumption,
 )
 from pytensor.tensor.blockwise import Blockwise
 
 
+@register_universal_assumption(Blockwise)
 def _blockwise_delegate(key, op, feature, fgraph, node, input_states):
     """Delegate assumption inference to the ``core_op`` of a Blockwise wrapper."""
     return infer_assumption_for_node(
         key, op.core_op, feature, fgraph, node, input_states
     )
-
-
-for _key in ALL_KEYS:
-    register_assumption(_key, Blockwise)(_blockwise_delegate)
diff --git a/pytensor/assumptions/bundles.py b/pytensor/assumptions/bundles.py
new file mode 100644
index 0000000000..001eaf4bf5
--- /dev/null
+++ b/pytensor/assumptions/bundles.py
@@ -0,0 +1,55 @@
+from pytensor.assumptions.alloc import alloc_propagates_matrix_property
+from pytensor.assumptions.core import AssumptionKey, register_assumption
+from pytensor.assumptions.dimshuffle import dimshuffle_propagates_matrix_property
+from pytensor.assumptions.reshape import (
+    join_dims_propagates_matrix_property,
+    split_dims_propagates_matrix_property,
+)
+from pytensor.assumptions.shape import (
+    reshape_propagates_matrix_property,
+    specify_shape_propagates_matrix_property,
+)
+from pytensor.assumptions.subtensor import (
+    incsubtensor_propagates_matrix_property,
+    subtensor_propagates_matrix_property,
+)
+from pytensor.tensor.basic import Alloc
+from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.reshape import JoinDims, SplitDims
+from pytensor.tensor.shape import Reshape, SpecifyShape
+from pytensor.tensor.subtensor import IncSubtensor, Subtensor
+
+
+def register_matrix_property_rules(key: AssumptionKey) -> None:
+    """Register the standard propagation rules for a property of the trailing two axes.
+
+    Every rule here answers one question: does the Op leave the trailing two axes
+    undisturbed? The bundle thus suits any property of a matrix that batch dimensions
+    carry elementwise, such as triangularity or a fixed sparsity pattern.
+
+    Rules are tried in registration order until one returns a non-UNKNOWN state, so a
+    key needing different behavior for one Op registers its own with
+    ``register_assumption(..., prepend=True)``.
+
+    Parameters
+    ----------
+    key : AssumptionKey
+        The property to install the rules for.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from pytensor.assumptions import AssumptionKey, register_matrix_property_rules
+
+        TOEPLITZ = AssumptionKey("toeplitz", short_name="toep")
+        register_matrix_property_rules(TOEPLITZ)
+    """
+    register_assumption(key, DimShuffle)(dimshuffle_propagates_matrix_property)
+    register_assumption(key, Reshape)(reshape_propagates_matrix_property)
+    register_assumption(key, SpecifyShape)(specify_shape_propagates_matrix_property)
+    register_assumption(key, JoinDims)(join_dims_propagates_matrix_property)
+    register_assumption(key, SplitDims)(split_dims_propagates_matrix_property)
+    register_assumption(key, Alloc)(alloc_propagates_matrix_property)
+    register_assumption(key, Subtensor)(subtensor_propagates_matrix_property)
+    register_assumption(key, IncSubtensor)(incsubtensor_propagates_matrix_property)
diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py
index 9b1195fc5d..9d70d92407 100644
--- a/pytensor/assumptions/core.py
+++ b/pytensor/assumptions/core.py
@@ -1,12 +1,14 @@
 from collections import deque
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from enum import IntFlag, auto
 from typing import Any
 
 from pytensor.graph import Apply, FunctionGraph, Op
+from pytensor.graph.basic import Variable
 from pytensor.graph.features import AlreadyThere, Feature
 from pytensor.graph.traversal import walk_toposort
+from pytensor.tensor import TensorLike
 from pytensor.tensor.variable import TensorConstant
 
 
@@ -41,16 +43,107 @@ def join(cls, left: "FactState", right: "FactState") -> "FactState":
 class AssumptionKey:
     """Identifies a named structural property (e.g. "diagonal" or "triangular").
 
-    ``short_name`` is an abbreviated label used by ``debugprint(print_assumptions=True)``;
-    it falls back to ``name`` when empty.
+    Constructing a key registers it in :data:`KEY_REGISTRY` and installs every rule
+    declared with :func:`register_universal_assumption` for it, which is all a
+    downstream library must do to add a property of its own.
+
+    ``name`` is the key's identity: it is what :func:`assume` accepts as a keyword and
+    what two keys may not share unless they are identical in every field.
+
+    Parameters
+    ----------
+    name : str
+        Unique identifier for the property, also the keyword :func:`assume` accepts.
+    short_name : str, optional
+        Abbreviated label used by ``debugprint(print_assumptions=True)``. Falls back
+        to ``name`` when empty.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        import pytensor.tensor as pt
+        from pytensor.assumptions import AssumptionKey
+
+        TOEPLITZ = AssumptionKey("toeplitz", short_name="toep")
+
+        x = TOEPLITZ.assume(pt.matrix("x"))
+        TOEPLITZ.holds(x)  # True
+        TOEPLITZ.holds(pt.exp(x))  # False -- no rule teaches it about Elemwise yet
     """
 
     name: str
     short_name: str = ""
 
+    def __post_init__(self) -> None:
+        registered = KEY_REGISTRY.get(self.name)
+        if registered is not None:
+            if registered != self:
+                raise ValueError(
+                    f"An assumption named {self.name!r} is already registered as "
+                    f"{registered!r} with different metadata. Assumption names are "
+                    f"global identifiers; pick a distinct one."
+                )
+            # Re-creating an identical key (a module imported twice) is a no-op:
+            # equal keys hash alike, so the installed rules already apply to it.
+            return
+
+        KEY_REGISTRY[self.name] = self
+        for op_types, fn in UNIVERSAL_RULES:
+            _install_rule(self, op_types, fn)
+
+    def __reduce__(self):
+        # Unpickle through the constructor: the default dataclass path skips __init__,
+        # leaving a key that is in no registry and has no rules installed -- not even
+        # the one that reads declarations back off SpecifyAssumptions.
+        return type(self), (self.name, self.short_name)
+
     def __repr__(self) -> str:
         return self.name
 
+    def assume(self, x: TensorLike, *, state: bool = True):
+        """Return a view of *x* declaring this assumption.
+
+        Parameters
+        ----------
+        x : tensor-like
+            The input to annotate.
+        state : bool, optional
+            Whether to assert the property holds or that it does not. Default True.
+
+        Returns
+        -------
+        out : TensorVariable
+            A view of *x* with the assumption attached.
+        """
+        from pytensor.assumptions.specify import SpecifyAssumptions
+
+        fact = FactState.TRUE if state else FactState.FALSE
+        return SpecifyAssumptions({self: fact})(x)
+
+    def holds(self, var: Variable, fgraph: FunctionGraph | None = None) -> bool:
+        """Return True iff this assumption is provably TRUE for *var*.
+
+        Parameters
+        ----------
+        var : Variable
+            The variable to ask about.
+        fgraph : FunctionGraph, optional
+            Graph to resolve the question in. Pass one when asking about several
+            variables of the same graph: without it a throwaway ``FunctionGraph`` is
+            built per call, which walks the ancestors of *var* and discards the
+            inference cache afterwards.
+
+        Returns
+        -------
+        holds : bool
+            True when the property is known to hold, False when it is known not to
+            hold or is simply unknown.
+        """
+        if fgraph is None:
+            fgraph = FunctionGraph(outputs=[var], clone=False)
+        return check_assumption(fgraph, var, self)
+
 
 class ConflictingAssumptionsError(ValueError):
     """Raised when joining evidence about a (variable, key) produces ``FactState.CONFLICT``.
@@ -72,6 +165,73 @@ class ConflictingAssumptionsError(ValueError):
 # Rules are tried in registration order; the first to return TRUE wins.
 ASSUMPTION_INFER_REGISTRY: dict[tuple[AssumptionKey, type], list[InferFactFn]] = {}
 
+# Every AssumptionKey ever constructed, by name. Downstream libraries join the system
+# by constructing a key; nothing else is required of them. This resolves the names
+# :func:`assume` takes as keywords -- graphs themselves carry keys, not names, so
+# nothing downstream of graph construction needs to look anything up here.
+KEY_REGISTRY: dict[str, AssumptionKey] = {}
+
+# Rules that hold for every key regardless of what the key means, as (op_types, fn)
+# pairs. Kept separately from ASSUMPTION_INFER_REGISTRY so that keys created *after*
+# the rule is declared still receive it -- see AssumptionKey.__post_init__.
+UNIVERSAL_RULES: list[tuple[tuple[type, ...], InferFactFn]] = []
+
+
+def _install_rule(
+    key: AssumptionKey, op_types: tuple[type, ...], fn: InferFactFn
+) -> None:
+    for op_type in op_types:
+        ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []).append(fn)
+
+
+def register_universal_assumption(
+    *op_types: type,
+) -> Callable[[InferFactFn], InferFactFn]:
+    """Decorator registering an inference rule that applies to *every* assumption key.
+
+    Use this for rules that are indifferent to what the property means -- an Op that
+    forwards its input unchanged, or one that delegates to another Op. The rule is
+    installed for keys that already exist and for every key created later.
+
+    Parameters
+    ----------
+    *op_types : type
+        Op classes the rule applies to.
+    """
+
+    def decorator(fn: InferFactFn) -> InferFactFn:
+        UNIVERSAL_RULES.append((op_types, fn))
+        for key in KEY_REGISTRY.values():
+            _install_rule(key, op_types, fn)
+        return fn
+
+    return decorator
+
+
+class KeyRegistryView:
+    """Live, read-only view of every registered :class:`AssumptionKey`."""
+
+    __slots__ = ()
+
+    def __iter__(self) -> Iterator[AssumptionKey]:
+        # Snapshot: a rule that constructs a key would otherwise resize the registry
+        # mid-iteration.
+        return iter(tuple(KEY_REGISTRY.values()))
+
+    def __contains__(self, value: object) -> bool:
+        # Only a key can be registered, so the isinstance both guards the ``name``
+        # access and lets the lookup be a hit rather than a scan.
+        return (
+            isinstance(value, AssumptionKey) and KEY_REGISTRY.get(value.name) == value
+        )
+
+    def __len__(self) -> int:
+        return len(KEY_REGISTRY)
+
+    def __repr__(self) -> str:
+        return f"({', '.join(KEY_REGISTRY)})"
+
+
 # Registry mapping assumptions to other assumptions they imply.  For example, a "diagonal" matrix is also "symmetric"
 # and "triangular".  This is consulted after all other inference rules to derive additional facts.
 IMPLIES: dict[AssumptionKey, list[AssumptionKey]] = {}
@@ -116,7 +276,10 @@ def register_constant_inference(key: AssumptionKey, fn: ConstantInferFn) -> None
     PERMUTATION,
 )
 
-ALL_KEYS = (*MATRIX_KEYS, UNIQUE_INDICES)
+# Live view rather than a tuple: a key registered by a downstream library shows up here
+# too, so anything that iterates the keys at call time (debugprint, the drain rewrite)
+# covers it without further registration.
+ALL_KEYS = KeyRegistryView()
 
 # Implications about structural properties derivably from other structural properties
 register_implies(DIAGONAL, LOWER_TRIANGULAR, UPPER_TRIANGULAR, SYMMETRIC)
@@ -125,17 +288,32 @@ def register_constant_inference(key: AssumptionKey, fn: ConstantInferFn) -> None
 
 
 def register_assumption(
-    key: AssumptionKey, *op_types: type
+    key: AssumptionKey, *op_types: type, prepend: bool = False
 ) -> Callable[[InferFactFn], InferFactFn]:
     """Decorator that registers an inference rule for ``(key, op_type)`` pairs.
 
     The decorated function is called as ``fn(key, op, feature, fgraph, node, input_states)``
     and must return a list of :class:`FactState` with one entry per node output.
+
+    Parameters
+    ----------
+    key : AssumptionKey
+        The property the rule infers.
+    *op_types : type
+        Op classes the rule applies to.
+    prepend : bool, optional
+        Run this rule ahead of those already registered for the same pair rather than
+        after them. Rules are tried in order until one returns a non-UNKNOWN state, so
+        this is how a key overrides a rule it inherited from a bundle. Default False.
     """
 
     def decorator(fn: InferFactFn) -> InferFactFn:
         for op_type in op_types:
-            ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []).append(fn)
+            rules = ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), [])
+            if prepend:
+                rules.insert(0, fn)
+            else:
+                rules.append(fn)
         return fn
 
     return decorator
diff --git a/pytensor/assumptions/dimshuffle.py b/pytensor/assumptions/dimshuffle.py
index e60edcd219..0f622e73a1 100644
--- a/pytensor/assumptions/dimshuffle.py
+++ b/pytensor/assumptions/dimshuffle.py
@@ -1,4 +1,4 @@
-from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption
+from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption
 from pytensor.tensor.elemwise import DimShuffle
 
 
@@ -29,5 +29,5 @@ def dimshuffle_propagates_matrix_property(
     return [FactState.UNKNOWN]
 
 
-for _key in ALL_KEYS:
+for _key in MATRIX_KEYS:
     register_assumption(_key, DimShuffle)(dimshuffle_propagates_matrix_property)
diff --git a/pytensor/assumptions/reshape.py b/pytensor/assumptions/reshape.py
index 53ec7d3a57..41a5f09a34 100644
--- a/pytensor/assumptions/reshape.py
+++ b/pytensor/assumptions/reshape.py
@@ -1,4 +1,4 @@
-from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption
+from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption
 from pytensor.tensor.reshape import JoinDims, SplitDims
 
 
@@ -23,6 +23,6 @@ def split_dims_propagates_matrix_property(
     return [FactState.UNKNOWN]
 
 
-for _key in ALL_KEYS:
+for _key in MATRIX_KEYS:
     register_assumption(_key, JoinDims)(join_dims_propagates_matrix_property)
     register_assumption(_key, SplitDims)(split_dims_propagates_matrix_property)
diff --git a/pytensor/assumptions/shape.py b/pytensor/assumptions/shape.py
index 1e907d94a5..a036de18b3 100644
--- a/pytensor/assumptions/shape.py
+++ b/pytensor/assumptions/shape.py
@@ -1,4 +1,4 @@
-from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption
+from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption
 from pytensor.tensor.shape import Reshape, SpecifyShape
 
 
@@ -25,6 +25,6 @@ def reshape_propagates_matrix_property(
     return [input_states[0]]
 
 
-for _key in ALL_KEYS:
+for _key in MATRIX_KEYS:
     register_assumption(_key, SpecifyShape)(specify_shape_propagates_matrix_property)
     register_assumption(_key, Reshape)(reshape_propagates_matrix_property)
diff --git a/pytensor/assumptions/specify.py b/pytensor/assumptions/specify.py
index bb9f77d6ae..a79ce371ff 100644
--- a/pytensor/assumptions/specify.py
+++ b/pytensor/assumptions/specify.py
@@ -1,6 +1,11 @@
 from collections.abc import Sequence
 
-from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption
+from pytensor.assumptions.core import (
+    KEY_REGISTRY,
+    AssumptionKey,
+    FactState,
+    register_universal_assumption,
+)
 from pytensor.compile.ops import TypeCastingOp
 from pytensor.graph.basic import Apply, Variable
 from pytensor.tensor import TensorLike
@@ -10,26 +15,42 @@
 class SpecifyAssumptions(TypeCastingOp):
     """No-op that declares structural assumptions on a tensor for use by graph rewrites.
 
-    ``assumptions`` is a tuple of ``(name, FactState)`` pairs sorted by ``name``, where
-    ``name`` matches the name of an :class:`AssumptionKey`. Two instances with the same
-    fact set compare equal via ``__props__``, so PyTensor's graph merge collapses
-    duplicates.
+    ``assumptions`` is a tuple of ``(AssumptionKey, FactState)`` pairs sorted by key
+    name. Declaring a fact therefore requires holding the key itself, and constructing
+    a key registers it, so a graph cannot carry an assumption the system has never
+    heard of. Two instances with the same fact set compare equal via ``__props__``, so
+    PyTensor's graph merge collapses duplicates.
+
+    Parameters
+    ----------
+    assumptions : dict mapping AssumptionKey to FactState
+        The facts to declare.
     """
 
     __props__ = ("assumptions",)
 
-    assumptions: tuple[tuple[str, FactState], ...]
+    assumptions: tuple[tuple[AssumptionKey, FactState], ...]
 
-    def __init__(self, assumptions: dict[str, FactState]):
+    def __init__(self, assumptions: dict[AssumptionKey, FactState]):
         super().__init__()
+        passed_by_name = [
+            key for key in assumptions if not isinstance(key, AssumptionKey)
+        ]
+        if passed_by_name:
+            raise TypeError(
+                f"SpecifyAssumptions is keyed by AssumptionKey, not by name: "
+                f"{passed_by_name!r}. Pass the key objects, or declare by name "
+                f"with assume()."
+            )
         self.assumptions = tuple(
-            (name, FactState(state)) for name, state in sorted(assumptions.items())
+            (key, FactState(state))
+            for key, state in sorted(assumptions.items(), key=lambda kv: kv[0].name)
         )
 
     def __str__(self):
         facts = ", ".join(
-            name if state is FactState.TRUE else f"!{name}"
-            for name, state in self.assumptions
+            key.name if state is FactState.TRUE else f"!{key.name}"
+            for key, state in self.assumptions
         )
         return f"{type(self).__name__}{{{facts}}}"
 
@@ -48,13 +69,14 @@ def pullback(
         return list(output_cotangents)
 
 
+@register_universal_assumption(SpecifyAssumptions)
 def specify_assumption_rule(key, op, feature, fgraph, node, input_states):
     """Report the declared state for ``key`` joined with whatever inference derived
     from the input. The join surfaces ``ConflictingAssumptionsError`` when the user
     asserts a state that contradicts what the system can prove (e.g. asserting
     ``diagonal=False`` on something proved diagonal)."""
-    for name, state in op.assumptions:
-        if name == key.name:
+    for declared_key, state in op.assumptions:
+        if declared_key == key:
             return [FactState.join(state, input_states[0])]
     return [input_states[0]]
 
@@ -70,6 +92,7 @@ def assume(
     selection: bool | None = None,
     permutation: bool | None = None,
     unique_indices: bool | None = None,
+    **assumptions: bool | None,
 ):
     """Attach structural assumptions to a symbolic tensor.
 
@@ -103,6 +126,9 @@ def assume(
         aliases a non-negative one (e.g. ``-1`` and ``n-1``). Such an index can
         never enlarge the axis it indexes, so it can be lifted earlier through
         operations without risk of duplicating computation.
+    **assumptions : bool, optional
+        Assumptions registered by downstream libraries, passed by key name, e.g.
+        ``time_varying=True``.
 
     Returns
     -------
@@ -111,15 +137,18 @@ def assume(
 
     Examples
     --------
-    >>> import pytensor.tensor as pt
-    >>> x = pt.dmatrix("x")
-    >>> x_diag = assume(x, diagonal=True)
-    >>> x_not_sym = assume(x, symmetric=False)
+    .. code-block:: python
+
+        import pytensor.tensor as pt
+
+        x = pt.dmatrix("x")
+        x_diag = assume(x, diagonal=True)
+        x_not_sym = assume(x, symmetric=False)
     """
     if not isinstance(x, Variable):
         x = as_tensor_variable(x)
 
-    values = {
+    core_values = {
         "diagonal": diagonal,
         "lower_triangular": lower_triangular,
         "upper_triangular": upper_triangular,
@@ -130,17 +159,23 @@ def assume(
         "permutation": permutation,
         "unique_indices": unique_indices,
     }
-    assumptions = {
-        name: FactState.TRUE if value else FactState.FALSE
-        for name, value in values.items()
+
+    unknown = [name for name in assumptions if name not in KEY_REGISTRY]
+    if unknown:
+        extensions = sorted(KEY_REGISTRY.keys() - core_values.keys())
+        raise ValueError(
+            f"Unknown assumption(s): {', '.join(unknown)}. Registered extension "
+            f"assumptions are: {', '.join(extensions) if extensions else '(none)'}. "
+            f"Register a new one by constructing an AssumptionKey."
+        )
+
+    declared = {
+        KEY_REGISTRY[name]: FactState.TRUE if value else FactState.FALSE
+        for name, value in (core_values | assumptions).items()
         if value is not None
     }
 
-    if not assumptions:
+    if not declared:
         return x
 
-    return SpecifyAssumptions(assumptions)(x)
-
-
-for _key in ALL_KEYS:
-    register_assumption(_key, SpecifyAssumptions)(specify_assumption_rule)
+    return SpecifyAssumptions(declared)(x)
diff --git a/pytensor/assumptions/subtensor.py b/pytensor/assumptions/subtensor.py
index 8f13338676..b7ec162c1b 100644
--- a/pytensor/assumptions/subtensor.py
+++ b/pytensor/assumptions/subtensor.py
@@ -1,7 +1,7 @@
 from pytensor.assumptions.core import (
-    ALL_KEYS,
     DIAGONAL,
     LOWER_TRIANGULAR,
+    MATRIX_KEYS,
     POSITIVE_DEFINITE,
     SYMMETRIC,
     UPPER_TRIANGULAR,
@@ -82,5 +82,5 @@ def incsubtensor_propagates_matrix_property(
     return true_if(base_state is FactState.TRUE and value_state is FactState.TRUE)
 
 
-for _key in ALL_KEYS:
+for _key in MATRIX_KEYS:
     register_assumption(_key, IncSubtensor)(incsubtensor_propagates_matrix_property)
diff --git a/pytensor/tensor/rewriting/assumptions.py b/pytensor/tensor/rewriting/assumptions.py
index a60b7a798d..3740863eef 100644
--- a/pytensor/tensor/rewriting/assumptions.py
+++ b/pytensor/tensor/rewriting/assumptions.py
@@ -1,12 +1,9 @@
-from pytensor.assumptions import ALL_KEYS, AssumptionFeature
+from pytensor.assumptions import AssumptionFeature
 from pytensor.assumptions.specify import SpecifyAssumptions
 from pytensor.compile.mode import optdb
 from pytensor.graph.rewriting.basic import GraphRewriter
 
 
-_KEY_BY_NAME = {key.name: key for key in ALL_KEYS}
-
-
 class DrainSpecifyAssumptions(GraphRewriter):
     """Drain ``SpecifyAssumptions`` declarations into the ``AssumptionFeature`` and
     remove the marker nodes.
@@ -38,16 +35,15 @@ def apply(self, fgraph):
         for node in nodes:
             [out] = node.outputs
             # Resolve the asserted facts into the cache.
-            for name, _ in node.op.assumptions:
-                assumption_feature.get(out, _KEY_BY_NAME[name])
-            # Drain the marker: redirect its consumers to the raw input,
-            # peeling nested SpecifyAssumptions so a single replace_all
-            # collapses ``assume(assume(...))`` chains all the way down.
+            for key, _ in node.op.assumptions:
+                assumption_feature.get(out, key)
+            # Drain the marker: redirect its consumers to the raw input, peeling
+            # already-drained nested SpecifyAssumptions -- ``nodes`` is in toposort
+            # order, so a single replace_all collapses ``assume(assume(...))`` chains
+            # all the way down.
             inp = node.inputs[0]
-            while inp.owner is not None and isinstance(
-                inp.owner.op, SpecifyAssumptions
-            ):
-                inp = inp.owner.inputs[0]
+            while inp in replacements:
+                inp = replacements[inp]
             replacements[out] = inp
 
         fgraph.replace_all(
diff --git a/tests/assumptions/conftest.py b/tests/assumptions/conftest.py
index b82e5d10c3..a8bd3d1990 100644
--- a/tests/assumptions/conftest.py
+++ b/tests/assumptions/conftest.py
@@ -12,6 +12,8 @@ def _snapshot_assumption_registries():
     """Restore module-global assumption registries after each test."""
     infer_snapshot = copy.deepcopy(_assumptions_core.ASSUMPTION_INFER_REGISTRY)
     implies_snapshot = copy.deepcopy(_assumptions_core.IMPLIES)
+    key_snapshot = dict(_assumptions_core.KEY_REGISTRY)
+    universal_snapshot = list(_assumptions_core.UNIVERSAL_RULES)
     try:
         yield
     finally:
@@ -19,6 +21,9 @@ def _snapshot_assumption_registries():
         _assumptions_core.ASSUMPTION_INFER_REGISTRY.update(infer_snapshot)
         _assumptions_core.IMPLIES.clear()
         _assumptions_core.IMPLIES.update(implies_snapshot)
+        _assumptions_core.KEY_REGISTRY.clear()
+        _assumptions_core.KEY_REGISTRY.update(key_snapshot)
+        _assumptions_core.UNIVERSAL_RULES[:] = universal_snapshot
 
 
 def make_fgraph(*outputs, **kwargs):
diff --git a/tests/assumptions/test_alloc.py b/tests/assumptions/test_alloc.py
index c248770800..596ded986b 100644
--- a/tests/assumptions/test_alloc.py
+++ b/tests/assumptions/test_alloc.py
@@ -10,6 +10,7 @@
     PERMUTATION,
     POSITIVE_DEFINITE,
     SYMMETRIC,
+    UNIQUE_INDICES,
     UPPER_TRIANGULAR,
     FactState,
 )
@@ -145,3 +146,19 @@ def test_alloc_broadcast_vector_value_is_unknown():
     y = pt.alloc(v, 4, 4)
     _, af = make_fgraph(y)
     assert af.get(y, SYMMETRIC) == FactState.UNKNOWN
+
+
+def test_unique_indices_survives_no_broadcast():
+    """Alloc repeats entries, so a uniqueness claim must not carry through it.
+
+    The matrix-property rules propagate anything whose trailing two axes are untouched,
+    which is wrong for a claim about the values themselves.
+    """
+    idx = pt.matrix("idx", shape=(2, 3), dtype="int64")
+    broadcast = pt.alloc(assume(idx, unique_indices=True), 4, 2, 3)
+
+    _, af = make_fgraph(broadcast)
+    assert af.get(broadcast, UNIQUE_INDICES) is not FactState.TRUE
+
+    repeated = broadcast.eval({idx: np.arange(6).reshape(2, 3)})
+    assert len(np.unique(repeated)) < repeated.size
diff --git a/tests/assumptions/test_bundles.py b/tests/assumptions/test_bundles.py
new file mode 100644
index 0000000000..9cd210a3cc
--- /dev/null
+++ b/tests/assumptions/test_bundles.py
@@ -0,0 +1,86 @@
+import pytest
+
+import pytensor.tensor as pt
+from pytensor.assumptions import (
+    MATRIX_KEYS,
+    AssumptionKey,
+    FactState,
+    assume,
+    register_assumption,
+    register_matrix_property_rules,
+)
+from pytensor.assumptions.core import ASSUMPTION_INFER_REGISTRY
+from pytensor.tensor.basic import Alloc
+from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.reshape import JoinDims, SplitDims
+from pytensor.tensor.shape import Reshape, SpecifyShape
+from pytensor.tensor.subtensor import IncSubtensor
+from tests.assumptions.conftest import make_fgraph
+
+
+def test_bundle_propagates_through_the_standard_ops():
+    TOEPLITZ = AssumptionKey("toeplitz", short_name="toep")
+    register_matrix_property_rules(TOEPLITZ)
+
+    x = assume(pt.tensor("x", shape=(4, 3, 3)), toeplitz=True)
+
+    indexed = x[0]
+    reshaped = x.reshape((2, 2, 3, 3))
+    shape_specified = pt.specify_shape(x, (4, 3, 3))
+
+    _, af = make_fgraph(indexed, reshaped, shape_specified)
+    assert af.check(indexed, TOEPLITZ)
+    assert af.check(reshaped, TOEPLITZ)
+    assert af.check(shape_specified, TOEPLITZ)
+
+
+def test_bundle_leaves_the_core_axes_alone():
+    """The rules protect the trailing two axes -- disturbing them stops propagation."""
+    TOEPLITZ = AssumptionKey("toeplitz", short_name="toep")
+    register_matrix_property_rules(TOEPLITZ)
+
+    x = assume(pt.matrix("x", shape=(3, 3)), toeplitz=True)
+    transposed = x.T
+
+    _, af = make_fgraph(transposed)
+    assert af.get(transposed, TOEPLITZ) is FactState.UNKNOWN
+
+
+@pytest.mark.parametrize(
+    "prepend, expected",
+    [(True, FactState.FALSE), (False, FactState.TRUE)],
+    ids=["prepend-wins", "append-loses"],
+)
+def test_prepend_decides_which_rule_answers(prepend, expected):
+    """Expand-dims is a case the bundle answers, so only order decides the outcome."""
+    TOEPLITZ = AssumptionKey("toeplitz", short_name="toep")
+    register_matrix_property_rules(TOEPLITZ)
+
+    @register_assumption(TOEPLITZ, DimShuffle, prepend=prepend)
+    def _never_survives_expand_dims(key, op, feature, fgraph, node, input_states):
+        return [FactState.FALSE] if op.is_expand_dims else [FactState.UNKNOWN]
+
+    x = assume(pt.matrix("x", shape=(3, 3)), toeplitz=True)
+    expanded = x[None]
+
+    _, af = make_fgraph(expanded)
+    assert af.get(expanded, TOEPLITZ) is expected
+
+
+@pytest.mark.parametrize("key", MATRIX_KEYS, ids=lambda k: k.name)
+@pytest.mark.parametrize(
+    "op_type",
+    [DimShuffle, Reshape, SpecifyShape, JoinDims, SplitDims, Alloc, IncSubtensor],
+    ids=lambda op: op.__name__,
+)
+def test_core_matrix_keys_carry_the_bundled_rules(key, op_type):
+    """The bundle stays in step with what the built-in matrix properties register.
+
+    ``Subtensor`` is excluded: ``SELECTION`` registers its own rule there instead of the
+    shared one, so that Op is deliberately not uniform across the built-in keys.
+    """
+    probe = AssumptionKey("probe")
+    register_matrix_property_rules(probe)
+
+    bundled = set(ASSUMPTION_INFER_REGISTRY[(probe, op_type)])
+    assert bundled <= set(ASSUMPTION_INFER_REGISTRY[(key, op_type)])
diff --git a/tests/assumptions/test_extension.py b/tests/assumptions/test_extension.py
new file mode 100644
index 0000000000..3c07a4dd71
--- /dev/null
+++ b/tests/assumptions/test_extension.py
@@ -0,0 +1,196 @@
+import pickle
+
+import pytest
+
+import pytensor.tensor as pt
+from pytensor.assumptions import (
+    ALL_KEYS,
+    KEY_REGISTRY,
+    SYMMETRIC,
+    AssumptionKey,
+    FactState,
+    register_assumption,
+    register_universal_assumption,
+)
+from pytensor.assumptions.core import ASSUMPTION_INFER_REGISTRY
+from pytensor.assumptions.specify import SpecifyAssumptions, assume
+from pytensor.printing import debugprint
+from pytensor.tensor.basic import AllocDiag, alloc_diag
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.rewriting.assumptions import DrainSpecifyAssumptions
+from tests.assumptions.conftest import make_fgraph
+
+
+def test_key_registers_itself():
+    key = AssumptionKey("time_varying", short_name="tv")
+    assert KEY_REGISTRY["time_varying"] is key
+    assert key in ALL_KEYS
+
+
+def test_identical_key_redefinition_is_idempotent():
+    """A module imported twice must not double-register its key or its rules."""
+    first = AssumptionKey("time_varying", short_name="tv")
+    n_keys = len(KEY_REGISTRY)
+    second = AssumptionKey("time_varying", short_name="tv")
+
+    assert second == first
+    assert len(KEY_REGISTRY) == n_keys
+    assert KEY_REGISTRY["time_varying"] is first
+
+
+def test_name_collision_with_different_metadata_raises():
+    AssumptionKey("time_varying", short_name="tv")
+    with pytest.raises(ValueError, match="already registered"):
+        AssumptionKey("time_varying", short_name="clashing")
+
+
+def test_assume_accepts_extension_key():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x = pt.tensor3("x")
+    x_tv = assume(x, time_varying=True)
+    _, af = make_fgraph(x_tv)
+    assert af.check(x_tv, TIME_VARYING)
+
+
+def test_assume_mixes_core_and_extension_keys():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x = pt.tensor("x", shape=(10, 3, 3))
+    x_both = assume(x, symmetric=True, time_varying=True)
+    _, af = make_fgraph(x_both)
+    assert af.check(x_both, SYMMETRIC)
+    assert af.check(x_both, TIME_VARYING)
+
+
+def test_assume_records_false_for_extension_key():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x = pt.tensor3("x")
+    x_static = assume(x, time_varying=False)
+    _, af = make_fgraph(x_static)
+    assert af.get(x_static, TIME_VARYING) is FactState.FALSE
+
+
+def test_assume_rejects_unregistered_name():
+    """A typo must not silently become a no-op, and the error must aid discovery."""
+    AssumptionKey("time_varying", short_name="tv")
+    x = pt.matrix("x")
+    with pytest.raises(ValueError, match="Unknown assumption\\(s\\): symmetrik"):
+        assume(x, symmetrik=True)
+    with pytest.raises(ValueError, match="are: time_varying"):
+        assume(x, symmetrik=True)
+
+
+def test_key_assume_and_holds():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x = pt.tensor3("x")
+
+    assert TIME_VARYING.holds(TIME_VARYING.assume(x))
+    assert not TIME_VARYING.holds(TIME_VARYING.assume(x, state=False))
+    assert not TIME_VARYING.holds(x)
+
+
+def test_holds_reuses_a_supplied_fgraph():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x_tv = TIME_VARYING.assume(pt.tensor3("x"))
+    y = pt.matrix("y")
+    fgraph, _ = make_fgraph(x_tv, y)
+
+    assert TIME_VARYING.holds(x_tv, fgraph)
+    assert not TIME_VARYING.holds(y, fgraph)
+
+
+def test_universal_rules_reach_a_later_key():
+    """A key created after ``blockwise`` was imported still gets its delegate."""
+    SPARSE = AssumptionKey("sparse")
+    register_assumption(SPARSE, AllocDiag)(
+        lambda key, op, feature, fgraph, node, input_states: [FactState.TRUE]
+    )
+
+    v_core = pt.vector("v", shape=(3,))
+    core_op = alloc_diag(v_core, offset=0, axis1=0, axis2=1).owner.op
+    v_batch = pt.matrix("v_batch", shape=(5, 3))
+    batched = Blockwise(core_op, signature="(n)->(n,n)")(v_batch)
+
+    _, af = make_fgraph(batched)
+    assert af.check(batched, SPARSE)
+
+
+def test_universal_rule_reaches_existing_keys():
+    """The decorator installs onto keys registered before it ran, not just after."""
+    EARLY = AssumptionKey("early")
+
+    @register_universal_assumption(AllocDiag)
+    def _always_true(key, op, feature, fgraph, node, input_states):
+        return [FactState.TRUE]
+
+    LATE = AssumptionKey("late")
+
+    diag = alloc_diag(pt.vector("v", shape=(3,)), offset=0, axis1=0, axis2=1)
+    _, af = make_fgraph(diag)
+
+    assert af.check(diag, EARLY)
+    assert af.check(diag, LATE)
+
+
+def test_membership_rejects_a_non_key_sharing_a_name():
+    """``in`` compares keys, not names -- a variable named after one is not a key."""
+    assert "symmetric" not in ALL_KEYS
+    assert pt.matrix("symmetric") not in ALL_KEYS
+
+
+def test_extension_key_appears_in_debugprint():
+    AssumptionKey("time_varying", short_name="tv")
+    x_tv = assume(pt.tensor3("x"), time_varying=True)
+    printed = debugprint(x_tv, print_assumptions=True, file="str")
+    assert "a={tv}" in printed
+
+
+def test_drain_resolves_extension_key():
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x = pt.tensor3("x")
+    fgraph, af = make_fgraph(assume(x, time_varying=True) + 1, inputs=[x])
+
+    DrainSpecifyAssumptions().apply(fgraph)
+
+    assert not any(
+        isinstance(node.op, SpecifyAssumptions) for node in fgraph.apply_nodes
+    )
+    assert af.check(x, TIME_VARYING)
+
+
+def test_graph_carries_keys_not_names():
+    """The declaration holds the key itself, so it cannot name an unregistered fact."""
+    TIME_VARYING = AssumptionKey("time_varying", short_name="tv")
+    x_tv = assume(pt.tensor3("x"), time_varying=True)
+    assert x_tv.owner.op.assumptions == ((TIME_VARYING, FactState.TRUE),)
+
+
+def test_declaring_by_name_is_rejected():
+    AssumptionKey("time_varying", short_name="tv")
+    with pytest.raises(TypeError, match="not by name"):
+        SpecifyAssumptions({"time_varying": FactState.TRUE})
+
+
+def test_key_survives_a_pickle_round_trip():
+    """A key restored from a cached graph re-registers and keeps its universal rules."""
+    key = AssumptionKey("time_varying", short_name="tv")
+    blob = pickle.dumps(key)
+    del KEY_REGISTRY["time_varying"]
+
+    restored = pickle.loads(blob)
+
+    assert KEY_REGISTRY["time_varying"] == restored
+    assert (restored, SpecifyAssumptions) in ASSUMPTION_INFER_REGISTRY
+    assert restored.holds(restored.assume(pt.tensor3("x")))
+
+
+def test_pickled_graph_keeps_its_declaration():
+    """A graph outliving the library that declared its key still drains the fact."""
+    key = AssumptionKey("time_varying", short_name="tv")
+    blob = pickle.dumps(key.assume(pt.tensor3("x")))
+    del KEY_REGISTRY["time_varying"]
+
+    restored_graph = pickle.loads(blob)
+    [(restored_key, state)] = restored_graph.owner.op.assumptions
+
+    assert state is FactState.TRUE
+    assert restored_key.holds(restored_graph)
diff --git a/tests/assumptions/test_specify.py b/tests/assumptions/test_specify.py
index 0847f48bd5..3a2a8ebd6d 100644
--- a/tests/assumptions/test_specify.py
+++ b/tests/assumptions/test_specify.py
@@ -68,8 +68,8 @@ def test_assume_chained_combines_facts():
 
 
 def test_specify_assumptions_op_equal_for_same_facts():
-    a = SpecifyAssumptions({"diagonal": FactState.TRUE, "symmetric": FactState.FALSE})
-    b = SpecifyAssumptions({"symmetric": FactState.FALSE, "diagonal": FactState.TRUE})
+    a = SpecifyAssumptions({DIAGONAL: FactState.TRUE, SYMMETRIC: FactState.FALSE})
+    b = SpecifyAssumptions({SYMMETRIC: FactState.FALSE, DIAGONAL: FactState.TRUE})
     assert a == b
     assert hash(a) == hash(b)