System for algebraic reasoning about linear alegbra by jessegrabowski · Pull Request #2032 · pymc-devs/pytensor

jessegrabowski · 2026-04-07T03:36:02Z

This PR is a proposal for a typing system for linear algebra primitives. The purpose is to enable graph-wide reasoning about the kinds of matrices, so that we can rewrite to efficient computational forms.

Current State

We currently have several linear algebra rewrites and plan to add more. These are tracked in #573. This is important because linear algebra is 1) ubiquotous, 2) expensive, and 3) inscrutiable. Pytensor's static graph representation and rewrite system is well positioned to provide users help writing the best possible programs involving heavy linear algebra, if only we can figure out what is going on.

Consider the motivating case of solve(A, b), when A = pt.diag(pt.arange(100_000)). This is an O(n^3) operation that will call out to specialized routines. But there is no need for this. Since A is diagonal, we can write this as an elementwise division b / pt.extract_diag(A).

How do we decide to do this? We:

Track the Solve Op
Check if either input "seems diagonal"
If so, do the rewrite

What seems diagonal? We assume an input is diagonal if it was created by pt.eye or pt.diag. Users cannot specify themselves whether input data is diagonal. If an Op get inbetween the known "diagonalish" Op and the Solve, we cannot detect diagonality. For example, we cannot rewrite solve(A * 3, b), because now the first input is Elemwise(Mul)(A, 3). Multiplication is diagonal-preserving (because it is zero-preserving), but since the known diagonal op is now buried inside the Elemwise, we're out of luck.

Proposal

My proposal is to reason about algebraic properties of matrices the same way we reason about shapes. For shapes, we attach a ShapeFeature to FunctionGraphs. Each Op has an infer_shape method that explains how the static shape propagates. Likewise, I propose an AssumptionFeature. Ops do not have infer_assumption methods. That would be too messy. Instead, we have a central ASSUMPTION_INFER_REGISTRY with keys (Op, AssumptionKey ) and values InferFactFn.

An AssumptionKey is just a marker class corresponding to an algebraic fact about a matrix, like DIAGONAL, LOWER_TRIANGULAR, ORTHOGONAL, POSITIVE, SEMIDEFINITE, and so on.
An InferFactFn has the following signature:

def infer_diagonal(op: Op, assumption_feature: AssumptionFeature, fgraph: FunctionGraph, node: Apply, input_facts: list[FactState]) -> list[FactState]:

Like other symbolic operations, the InferFactFunction takes an Op (plus global information about the graph it lives in, fgraph and assumption_feature), and information about its inputs (the list of FactStates) and returns a list of information about its outputs.

A FactState is a three-valued logic for assumption inference. The possible values are UNKNOWN, TRUE, or FALSE. A fourth state, CONFLICT, exists, but should never arise.

All facts about all Ops are assumed to be UNKNOWN unless we can prove otherwise. Proof comes from each Op's registered InferFactFunctions. The AssumptionFeature is responsible for gathering all the rules of fact propagation. An example of a simple fact is that all Eye Ops are DIAGONAL, provided it is 1) square and 2) offset of zero:

def true_if(cond: bool) -> list[FactState]:
    """``[TRUE]`` when *cond* holds, ``[UNKNOWN]`` otherwise."""
    return [FactState.TRUE] if cond else [FactState.UNKNOWN]

def eye_is_identity(node) -> bool:
    """True when an :class:`Eye` node produces the identity matrix (square, k == 0)."""
    n, m, k = node.inputs
    if not (isinstance(k, Constant) and k.data.item() == 0):
        return False
    if n is m:
        return True
    if isinstance(n, Constant) and isinstance(m, Constant):
        return n.data.item() == m.data.item()
    return False

@register_assumption(DIAGONAL, Eye)
def _eye(op, feature, fgraph, node, input_states):
    return true_if(eye_is_identity(node))

In a program, we can use the AssumptionFeature.get(x, AssumptionKey) to query about the state of a Variable. Here, we ask "is x diagonal?". Obviously it is:

import pytensor.tensor as pt
from pytensor.graph.fg import FunctionGraph
from pytensor.tensor.assumptions import AssumptionFeature, DIAGONAL
x = pt.eye(5)
fg = FunctionGraph([], [x])
af = AssumptionFeature()
fg.attach_feature(af)

af.get(x, DIAGONAL) # <FactState.TRUE: 1>

Where this becomes powerful is by accumulating InferFactFunctions. Many Ops preserve diagonality, like Cholesky or Inverse. We can use information about the inputs to these Ops to propogate fact information through the graph:

@register_assumption(DIAGONAL, Cholesky)
def _cholesky(op, feature, fgraph, node, input_states):
    return true_if(input_states[0])

@register_assumption(DIAGONAL, MatrixInverse)
def _inv(op, feature, fgraph, node, input_states):
    return true_if(input_states[0])

Now we don't lose information about x in deeper graphs, and are free to do more rewrites:

x = pt.eye(5)
y = pt.linalg.cholesky(x)
z = pt.linalg.inv(y)

fg = FunctionGraph([z], [x])
af = AssumptionFeature()
fg.attach_feature(af)

af.get(y, DIAGONAL) # <FactState.TRUE: 1>
af.get(z, DIAGONAL) # <FactState.TRUE: 1>

Of course can also reason conditionally. An IncSubtensor might be diagonal-preserving if we can prove that we're setting a value on the diagonal of the matrix. Otherwise we fall back to UNKNOWN:

from pytensor.tensor.subtensor import IncSubtensor

x = pt.eye(5)
i, j = pt.iscalars('i', 'j')
y = x[i, j].inc(3) # Cannot prove i != j at runtime, so UNKNOWN
z = x[i, i].inc(3) # i == i provable, so this is diagonal-preserving
fg = FunctionGraph([y, z], [x])
af = AssumptionFeature()
fg.attach_feature(af)

af.get(y, DIAGONAL) #<FactState.UNKNOWN: 0>
af.get(z, DIAGONAL) #<FactState.TRUE: 1>

Facts can also imply other facts. DIAGONAL matrices are also symmetrical. These general relationships can be registered and encoded as well. Continuing the example above:

from pytensor.tensor.assumptions import SYMMETRIC
af.get(z, SYMMETRIC) # <FactState.TRUE: 1>

Finally, users can specify facts about matrices using pt.specify_assumptions, the same way they are able to specify shapes.

x = pt.specify_assumptions(pt.dmatrix('x'), diagonal=True)
fg = FunctionGraph([x], [x])
af = AssumptionFeature()
fg.attach_feature(af)

af.get(x, DIAGONAL) # <FactState.TRUE: 1>

Benefits for rewriting:

FactStates are trivial to check. We can check any fact about any Variable in 5 lines:

def _is_diagonal(var, fgraph):
    """Check if *var* is diagonal using the AssumptionFeature."""

    af = getattr(fgraph, "assumption_feature", None)
    if af is None:
        af = AssumptionFeature()
        fgraph.attach_feature(af)

    return af.check(var, DIAGONAL)

We can reason globally about the graph
As noted above, information flows through the graph. As long as we have good coverage of fact rules for Ops, we can make statements about Variables at all levels of computation
InferFactFunctions are lightweight and easy to write
It is trivial to add new InferFactFunctions. LLMs can bang them out. They require only local reasoning about the Op and its immediate inputs.
FactStates allow non-trivial combinations of rewrites
One example I hit while working on this was a rewrite for DirectSolveLyapunov given diagonal inputs. Because there is a rule that kron preserves diagonality if both inputs are diagonal, the chain of rewrites from solve_discrete_lyapunov(diag(a), Q) -> Q.ravel() / (1 - outer(a, a).ravel()) is discovered by the rewrite system via:
rewrite_kron_diag_to_diag_outer: kron(diag(a), diag(b)) → diag(outer(a, b).ravel())
rewrite_solve_diag_to_division: solve(diag(x), b) → b / x

The key is that after the first rewrite produces a diagonal matrix (via alloc_diag), the assumption system recognizes it as diagonal (via AllocDiag being registered with DIAGONAL), and then the second rewrite kicks in.

Non-Goals

The purpose of this system is not to introduce an complete, closed algebra over all types. That is impossible. The goal is also not to complete the project of German romanticism. That is also impossible.

There should never be any "global state" (no Wolfram Assuming). Assumptions live on FunctionGraphs. There is never a need to to deal with global context, logical combinations, relational assumptions.
We are not trying to be smart or figure out as much as possible. On the contrary, we want to be very dumb. The Assumptions should be maximally conservative, and fall back to UNKNOWN whenever there is runtime ambiguity.
Assumption logic lives separately from all other machinery. No other part of Pytensor needs to know anything about them. Evaluation logic does not check assumptions. Perform methods don't dispatch on the FactState of the Node. We check during rewrites and that's it.
We do not aspire to provide Maple-style conditional dispatches. If we don't know, we just don't know. There is no LinearAlgebra[IsDefinite]. No symbolic conditionals as part of the core API.
We are not and cannot be a theorm prover. We do not present assumptions to the user as a tool for this. Assumptions are primarily an inner-api, and rewrite-focused.

ricardoV94 · 2026-04-07T06:56:26Z

I like the top message (didn't look at the code). Two notes:

1

Missing some discussion on how to preserve information across rewrites. ShapeFeature combines shape information when you replace a->b and you knew something of a and something of b (or just one of them).

This is specially relevant for constant folding (e.g, eye(int(1e6))) where it's much cheaper/obvious before vs after

But then, when do you decide to ask/start checking assumption? Because this can also be useless work (e.g. a graph without any linalg stuff).

Ordering is hard

2

Why do you assume that you can do all the reasoning you need to from op and inputs? It's a small note and I don't think a requirement/restriction in your proposal.

But eg checking for tridiagonal matrix creation (which we do) requires checking 2/3 nested set_subtensor nodes.

maresb · 2026-04-07T07:35:04Z

Very disappointed that German romanticism is out of scope

But seriously, this looks amazing. I'm not really capable of reviewing this but I am very excited for this one.

ricardoV94 · 2026-04-07T11:43:22Z

+    try:
+        val = get_underlying_scalar_constant_value(node.inputs[0])
+        if val == 0:
+            return [FactState.TRUE]


Do we want FactState=False? when it is not zero but still known constant?

ricardoV94 · 2026-04-07T11:45:52Z

+        return output_grads
+
+
+def specify_assumptions(


Shorter name? pytensor.assume?

(Also this need not be in tensor module, seems more generic that it)

I think it could be shorter too, the other arguments a user has to give to it I think will make it clear what it does, so the name can be on the short side. Outside of tensors/linear algebra stuff, what kinda examples are you thinking of @ricardoV94?

xtensor / scalar modules can also have assumptions, no reason it needs to be tied to tensor/

ricardoV94 · 2026-04-08T11:49:14Z

+    return true_if(eye_is_identity(node))
+
+
+@register_assumption(ORTHOGONAL, MatrixInverse)


Use BlockwiseOf?

I have a helper in assumptions/blockwise.py that pushes all the assumptions through blockwise. Thoughts?

lucianopaz

I really like the goal of the PR. I couldn't finish reviewing but I'll just leave the comments I've written down so far. I'll try to come back to this later and write something more decent.

lucianopaz · 2026-04-09T09:32:02Z

+) -> list[FactState]:
+    """Determine the *key* fact for every output of *node*.
+
+    Resolution order:


This is the opposite resolution order than the one used in blockwise. Which should be preferred?

in assumptions/blockwise.py or in the actual blockwise Op?

lucianopaz · 2026-04-09T09:40:46Z

+        state = FactState(state)
+        cache_key = (var, key)
+        old = self.user_facts.get(cache_key, FactState.UNKNOWN)
+        new = FactState.join(old, state)


If old is CONFLICT, wont join also return a conflict? It wont' overwrite the old fact. Is this intended?

Related to below, I also hate the CONFLICT state. If we switch to True / Unknown, this is moot. But in this case this was intended, I was thinking about conflict as "nan" behaves in float

lucianopaz · 2026-04-09T10:58:39Z

+
+@register_assumption(DIAGONAL, Dot)
+def _dot(op, feature, fgraph, node, input_states):
+    return true_if(all(input_states))


I just came back to this after having looked at the orthogonal assumptions module. How would you handle the case where Q @ Q.T = eye? In other words, if the two inputs are orthogonal and one is the transpose of the other. Their product would produce an identity matrix. Would you be able to get the assumption from a different set (orthogonality) while working through diagonality?

Good question. The system as it exists don't have good tooling for handling cross-facts like that. Need to pause and ponder.

In the latest commit I showed how this would work. The system can actually handle cross-facts, but it will be important to explain the rules:

You register fact functions according to the output. So for your example of Q @ Q.T = eye, we register it as a diagonal fact.

The fact function takes the feature as an input, so you can always query against other facts, like feature.check(Q, ORTHOGONAL)

So the fact function looks like:

@register_assumption(DIAGONAL, Dot) def _dot_orthogonal_xxt(op, feature, fgraph, node, input_states): """x @ x.T is diagonal (identity) when x is orthogonal.""" a, b = node.inputs b_owner = b.owner if ( feature.check(a, ORTHOGONAL) and b_owner is not None and isinstance(b_owner.op, DimShuffle) and b_owner.op.is_matrix_transpose and b_owner.inputs[0] is a ): return [FactState.TRUE] return [FactState.UNKNOWN]

zaxtax · 2026-04-15T12:25:31Z

This proposal looks great! I'm totally in favour. Related to non-goals, are there any formal properties we would like to maintain about our system? Obviously users can introduce invalid rewrites, but in terms of the core implementation, any invariants and properties we want to maintain?

jessegrabowski · 2026-04-18T23:18:51Z

Obviously users can introduce invalid rewrites, but in terms of the core implementation, any invariants and properties we want to maintain?

I think we want to be super defensive. I am already pushing it bit with this one:

pytensor/pytensor/tensor/assumptions/positive_definite.py

Line 120 in 4016443

def _dot(op, feature, fgraph, node, input_states):

I included it to be provocative and generate discussion. In most realistic cases it's true that X @ X.T will be positive definite, but we can never prove that X is full rank at graph construction time (well, unless it's data we can inspect...). We have a concept of "shape_unsafe" rewrites that might create invalid graphs if things shift underneath the rewrite. If we want to be more aggressive we would need to introduce something like that. My issue here is that "shape_unsafe" assumes that users know about shape unsafe and know to turn it off in what cases. I think that's a pretty heroic assumption.

`_is_provably_positive(var, strict=...)` proves `var > 0` (strict) or `var >= 0` (non-strict); `_is_provably_non_negative` becomes a thin `strict=False` wrapper. The `uint`-dtype, `Shape`/`Shape_i`, and `Cast` cases prove non-negativity but not strict positivity, so they are recognized only when `strict=False`. Results cache on `tag.is_positive` or `tag.is_non_negative` per the flag.

jessegrabowski · 2026-05-17T00:33:06Z

I added a print_assumptions option to dprint. Here's how it looks for a GP example:

 import pytensor
import pytensor.tensor as pt
from pytensor.assumptions import assume

X = pt.matrix("X")
ls = pt.scalar("ls")
eta = pt.scalar("eta")
v = pt.vector("v")

sq = pt.sum(X**2, axis=1)
D2 = sq[:, None] + sq[None, :] - 2 * (X @ X.T)

K = eta**2 * pt.exp(-0.5 * D2 / ls**2) + 1e-6 * pt.eye(X.shape[0])

# domain knowledge the structural engine can't derive: a valid kernel is PD
K = assume(K, positive_definite=True)

n = X.shape[0]
logp = (
    -0.5 * (v @ pt.linalg.inv(K) @ v)
    - 0.5 * pt.log(pt.linalg.det(K))
    - 0.5 * n * pt.log(2 * pt.pi)
)

f = pytensor.function([X, ls, eta, v], logp)
f.dprint(print_assumptions=True)

Output:

Composite{((-0.5 * i1) - i0)} [id A] 22
 ├─ Sum{axes=None} [id B] 20
 │  └─ Log [id C] 18
 │     └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id D] 15
 │        └─ Cholesky{lower=True, overwrite_a=True} [id E] a={tril} 13
 │           └─ Composite{((i0 * exp(((-0.5 * ((i2 + i3) - (2.0 * i4))) / i1))) + (1e-06 * i5))} [id F] 12
 │              ├─ Sqr [id G] a={sym} 9
 │              │  └─ ExpandDims{axes=(0, 1)} [id H] 4
 │              │     └─ eta [id I]
 │              ├─ Sqr [id J] a={sym} 8
 │              │  └─ ExpandDims{axes=(0, 1)} [id K] 3
 │              │     └─ ls [id L]
 │              ├─ ExpandDims{axis=1} [id M] 11
 │              │  └─ Sum{axis=1} [id N] 7
 │              │     └─ Sqr [id O] 2
 │              │        └─ X [id P]
 │              ├─ ExpandDims{axis=0} [id Q] 10
 │              │  └─ Sum{axis=1} [id N] 7
 │              │     └─ ···
 │              ├─ Dot [id R] a={pd} 6
 │              │  ├─ X [id P]
 │              │  └─ MatrixTranspose [id S] 'X.T' 1
 │              │     └─ X [id P]
 │              └─ Eye{dtype='float64'} [id T] a={diag, pd, orth} 5
 │                 ├─ Shape_i{0} [id U] 0
 │                 │  └─ X [id P]
 │                 ├─ Shape_i{0} [id U] 0
 │                 │  └─ ···
 │                 └─ 0 [id V]
 └─ Squeeze{axes=(0, 1)} [id W] 21
    └─ Dot [id X] 19
       ├─ ExpandDims{axis=0} [id Y] 17
       │  └─ SolveTriangular{unit_diagonal=False, lower=True, b_ndim=1, overwrite_b=False} [id Z] 14
       │     ├─ Cholesky{lower=True, overwrite_a=True} [id E] a={tril} 13
       │     │  └─ ···
       │     └─ v [id BA]
       └─ ExpandDims{axis=1} [id BB] 16
          └─ SolveTriangular{unit_diagonal=False, lower=True, b_ndim=1, overwrite_b=False} [id Z] 14
             └─ ···

Inner graphs:

Composite{((-0.5 * i1) - i0)} [id A]
 ← sub [id BC]
    ├─ mul [id BD]
    │  ├─ -0.5 [id BE]
    │  └─ i1 [id BF]
    └─ i0 [id BG]

Composite{((i0 * exp(((-0.5 * ((i2 + i3) - (2.0 * i4))) / i1))) + (1e-06 * i5))} [id F]
 ← add [id BH]
    ├─ mul [id BI]
    │  ├─ i0 [id BG]
    │  └─ exp [id BJ]
    │     └─ true_div [id BK]
    │        ├─ mul [id BL]
    │        │  ├─ -0.5 [id BM]
    │        │  └─ sub [id BN]
    │        │     ├─ add [id BO]
    │        │     │  ├─ i2 [id BP]
    │        │     │  └─ i3 [id BQ]
    │        │     └─ mul [id BR]
    │        │        ├─ 2.0 [id BS]
    │        │        └─ i4 [id BT]
    │        └─ i1 [id BF]
    └─ mul [id BU]
       ├─ 1e-06 [id BV]
       └─ i5 [id BW]

We can see the gram matrix (id R) was automatically tagged as positive definite a={pd}. This resulted in all the right rewrites firing -- inv -> solve(assume_a = pos) -> cholesky + chosolve. Since there was a cholesky laying around, we also got det(x) -> slogdet(chol(x)) -> log(diag(sum(chol(x)))). The engine doesn't know Mercer's theorem, so I had to tag K = posdef myself, otherwise we would have gotten generic solve and an LU decomposition instead (but that's still not too bad!)

jessegrabowski · 2026-05-17T06:12:12Z

I reorganized the tests in the way that I like (file test_assumption.py became folder tests/assumptions with one file per library file). There are a lot of small tests, but they all jointly run sub-second.

ricardoV94 · 2026-05-17T13:37:23Z

+def eye_identity_rule(key, op, feature, fgraph, node, input_states) -> list[FactState]:
+    """Rule body: TRUE when an :class:`Eye` node produces the identity matrix (square, k == 0)."""
+    n, m, k = node.inputs
+    if not (isinstance(k, TensorConstant) and k.data.item() == 0):


Similarly, if it's a constant you can prove True/False, only if not is it unknown?

Yes, but the case was a bit more subtle because this utility was used by diagonal, posdef, and triangular. The rule only returns False as written for posdef and diagonal, even if k !=0 we are stilll triangular. I split the helpers.

Per-Op inference rules for `DIAGONAL`, `LOWER_TRIANGULAR`, `UPPER_TRIANGULAR`, `SYMMETRIC`, `POSITIVE_DEFINITE`, and `ORTHOGONAL`, queryable from rewrites via `check_assumption(fgraph, var, key)`. Facts are inferred lazily and cached on a `FunctionGraph` feature attached on first use. `pt.assume(x, diagonal=True, ...)` asserts facts directly via a `SpecifyAssumptions` no-op wrapper. Implication graph wires `DIAGONAL ⇒ LOWER_/UPPER_TRIANGULAR, SYMMETRIC` and `POSITIVE_DEFINITE ⇒ SYMMETRIC`. Migrate `cholesky_ldotlt`, `cholesky_of_diag`, `inv_of_diag_to_diag_reciprocal`, `det_of_diag`, `psd_solve_to_chol_solve`, `useless_symmetric_transpose`, and `inv_to_solve` to `check_assumption`, keeping the original tag fallbacks alongside for PyMC compatibility. Add `paired_triangular_solves_to_cho_solve` to fuse the Cholesky→solve pair the migrated PSD path leaves behind. Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

`SpecifyShape`, `Reshape`, `JoinDims`, `SplitDims`, `Alloc`, `DimShuffle`, and `IncSubtensor` only rearrange batch axes when they leave the trailing two (core) axes intact. A per-matrix property then carries through, so each new rule forwards the input fact in that case and reports UNKNOWN otherwise. Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

jessegrabowski requested review from ricardoV94 and zaxtax April 7, 2026 03:36

jessegrabowski added feature request request discussion graph rewriting linalg Linear algebra labels Apr 7, 2026

jessegrabowski force-pushed the assumption-system branch from cece59a to 9f6e7f5 Compare April 7, 2026 03:48

jessegrabowski requested review from aseyboldt, bwengals and lucianopaz April 7, 2026 03:51

ricardoV94 reviewed Apr 7, 2026

View reviewed changes

ricardoV94 reviewed Apr 8, 2026

View reviewed changes

lucianopaz reviewed Apr 9, 2026

View reviewed changes

jessegrabowski force-pushed the assumption-system branch 2 times, most recently from 8ac904d to 5767b43 Compare April 11, 2026 21:40

jessegrabowski mentioned this pull request Apr 12, 2026

More rewrites for ExtractDiag #2045

Merged

3 tasks

jessegrabowski force-pushed the assumption-system branch 2 times, most recently from 1d87d8e to cc01e9d Compare April 17, 2026 03:42

jessegrabowski mentioned this pull request Apr 18, 2026

Add empty input guards to numba linalg #2057

Merged

jessegrabowski force-pushed the assumption-system branch 2 times, most recently from fa48d6d to 4676b07 Compare April 18, 2026 22:08

jessegrabowski force-pushed the assumption-system branch from 4016443 to a7394bd Compare April 18, 2026 23:41

jessegrabowski force-pushed the assumption-system branch 4 times, most recently from 769df10 to 54758e8 Compare May 2, 2026 16:54

jessegrabowski force-pushed the assumption-system branch from 7d6b5b3 to 95cceda Compare May 16, 2026 20:36

jessegrabowski force-pushed the assumption-system branch 5 times, most recently from d94d905 to 6d4187c Compare May 17, 2026 00:27

jessegrabowski force-pushed the assumption-system branch 2 times, most recently from d481ac9 to f4bb66f Compare May 17, 2026 06:10

jessegrabowski force-pushed the assumption-system branch from f4bb66f to 7e999ff Compare May 17, 2026 06:29

ricardoV94 reviewed May 17, 2026

View reviewed changes

Comment thread pytensor/assumptions/alloc.py Outdated

ricardoV94 reviewed May 17, 2026

View reviewed changes

Comment thread pytensor/assumptions/diagonal.py Outdated

ricardoV94 reviewed May 17, 2026

View reviewed changes

Comment thread pytensor/assumptions/diagonal.py

ricardoV94 force-pushed the assumption-system branch from 920e5d7 to adbcf32 Compare May 17, 2026 15:23

ricardoV94 reviewed May 17, 2026

View reviewed changes

Comment thread pytensor/assumptions/core.py Outdated

jessegrabowski force-pushed the assumption-system branch 2 times, most recently from f6aa532 to b4c1a3c Compare May 17, 2026 21:03

jessegrabowski and others added 8 commits May 17, 2026 16:32

Add print_assumptions argument to dprint

531a7fb

Add DrainSpecifyAssumptions rewrite

5aedb86

Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

Add diagonal-aware rewrites for dot, kron, expm

76cfdab

Add diagonal-aware decomposition and solver rewrites

24cdca5

Add structured-solve, eig->eigh, and det_of_triangular rewrites

4e83d33

Add orthogonal-aware inv, solve, and dot rewrites

6c8cb31

jessegrabowski force-pushed the assumption-system branch from 98e9b5a to 6c8cb31 Compare May 17, 2026 21:35

jessegrabowski merged commit 7d58ba6 into pymc-devs:main May 17, 2026
66 checks passed

		return true_if(eye_is_identity(node))


		@register_assumption(ORTHOGONAL, MatrixInverse)

Conversation

jessegrabowski commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Current State

Proposal

Benefits for rewriting:

Non-Goals

Uh oh!

ricardoV94 commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

1

2

Uh oh!

maresb commented Apr 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lucianopaz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax commented Apr 15, 2026

Uh oh!

jessegrabowski commented Apr 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jessegrabowski commented May 17, 2026

Uh oh!

jessegrabowski commented May 17, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

jessegrabowski commented Apr 7, 2026 •

edited

Loading

ricardoV94 commented Apr 7, 2026 •

edited

Loading

jessegrabowski commented Apr 18, 2026 •

edited

Loading