Have you ever had the chance to look at PyTensor and its mathematical rewrite machinery?
https://pytensor.readthedocs.io/
Obviously there are many systems out there to operate symbolically on mathematical expressions but because you’re here and we use it I couldn’t help but ask.
We do very analogous stuff (for different goals) like lifting indexing operations or partial constant folding of graphs. On the more bayesian side we make use of it for deriving probability densities/mass or marginalizing variables.