Skip to content

slogdet_specialization races factor-sharing rewrites inside the specialize equilibrium #2477

Description

@jessegrabowski

slogdet_specialization and the factorization-creating rewrites run in the same specialize equilibrium, and the winner on a det node depends on node visit order. When slogdet_specialization reaches the det before reuse_decomposition_multiple_solves has created the factor, the det becomes an SLogDet and det_of_matrix_factorized_elsewhere has nothing left to match, so the matrix is factored twice. Output order alone flips the outcome:

import pytensor
import pytensor.tensor as pt

A = pt.dmatrix("A")
logdet = pt.log(pt.abs(pt.linalg.det(A)))
b1, b2 = pt.dvector("b1"), pt.dvector("b2")
x1 = pt.linalg.solve(A, b1, b_ndim=1)
x2 = pt.linalg.solve(A, b2, b_ndim=1)

pytensor.function([A, b1, b2], [logdet, x1, x2])  # LUFactor and SLogDet: A factored twice
pytensor.function([A, b1, b2], [x1, x2, logdet])  # LUFactor only, det shares the factor

A fix that does not depend on ordering is the missing counterpart rewrite: track SLogDet and share an existing factorization of its input into (sign, logabsdet). Cholesky gives sign = 1 and logabsdet = 2 * sum(log(diag(L))); LU needs the pivot parity for the sign. Either firing order then converges to one factorization.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions