local_useless_unbatched_blockwise emits squeeze(expand_dims(x)) that nothing collapses, so factorizations stop merging

[jessegrabowski]

A Blockwise whose batch dims are all broadcastable is unbatched by local_useless_unbatched_blockwise into core_op(squeeze(expand_dims(x))), and the compiled graph keeps that Squeeze/ExpandDims pair. The rewrite is registered at optdb position 60, after every DimShuffle-collapsing pass, so merge3 at position 100 sees Cholesky(A) and Cholesky(squeeze(expand_dims(A))) as different nodes and a positive-definite matrix used both batched and unbatched is factored twice.

import pytensor
import pytensor.tensor as pt

A = pt.matrix("A", shape=(5, 5))
b = pt.tensor("b", shape=(3, 5, 1))

logdet = 2 * pt.log(pt.diagonal(pt.linalg.cholesky(A))).sum()
quad = (b * pt.linalg.solve(A, b, assume_a="pos", b_ndim=2)).sum()

pytensor.dprint(pytensor.function([A, b], logdet + quad))

Composite{((2.0 * i1) + i0)} [id A] 9
 ├─ FusedElemwise{Mul, reduce[add@(0, 1, 2)]} [id B] 8
 │  ├─ b [id C]
 │  └─ [Blockwise{CholeskySolve{lower=True, b_ndim=2, overwrite_b=False}, (m,m),(m,n)->(m,n)}] [id D] 7
 │     ├─ ExpandDims{axis=0} [id E] 6
 │     │  └─ Cholesky{lower=True, overwrite_a=False} [id F] 5
 │     │     └─ Squeeze{axis=0} [id G] 4
 │     │        └─ ExpandDims{axis=0} [id H] 3
 │     │           └─ A [id I]
 │     ├─ b [id C]
 │     └─ [5 1] [id J]
 └─ FusedElemwise{Log, reduce[add@(0,)]} [id K] 2
    └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id L] 1
       └─ Cholesky{lower=True, overwrite_a=False} [id M] 0
          └─ A [id I]

The same pair with no solve involved:

pytensor.dprint(pytensor.function([A], [pt.linalg.cholesky(A), pt.linalg.cholesky(A[None])]))

Cholesky{lower=True, overwrite_a=False} [id A] 4
 └─ A [id B]
ExpandDims{axis=0} [id C] 3
 └─ Cholesky{lower=True, overwrite_a=False} [id D] 2
    └─ Squeeze{axis=0} [id E] 1
       └─ ExpandDims{axis=0} [id F] 0
          └─ A [id B]

[N-0-MAD]

Hi @jessegrabowski I looked into this and was able to reproduce the issue.

The problem seems to be in local_useless_unbatched_blockwise: when removing the dummy batch dimensions, it creates Squeeze nodes but doesn't immediately apply the existing apply_local_dimshuffle_lift cleanup. This leaves the Squeeze(ExpandDims(x)) structure in the graph and prevents the equivalent factorizations from being merged.

I have a small fix that applies apply_local_dimshuffle_lift after the squeeze. I also added a regression test showing that the two equivalent Cholesky factorizations are reduced to a single Cholesky node.

I’ve tested it against the current implementation and the test passes; I also verified that the regression test fails on the pre-fix code (2 Cholesky nodes instead of 1).
I’ve pushed the commit here for you to take a look before I open a PR:
dee706d37 — Fix useless unbatched Blockwise rewrite

Would appreciate your thoughts on whether this is the direction you had in mind.

[ricardoV94]

@N-0-MAD the approach is more or less what I'd do yes