Add Feature that can go back and forward in rewrite history

[ricardoV94]

Description

I can imagine this being very powerful for debugging or reasoning about rewrites?

import pytensor
import pytensor.tensor as pt
from pytensor.graph.fg import FunctionGraph
from pytensor.graph.features import FullHistory
from pytensor.graph.rewriting.utils import rewrite_graph

x = pt.scalar("x")
out = pt.log(pt.exp(x) / pt.sum(pt.exp(x)))

fg = FunctionGraph(outputs=[out])
history = FullHistory()
fg.attach_feature(history)

rewrite_graph(fg, clone=False, include=("canonicalize", "stabilize"))

# Replay rewrites
history.start()
pytensor.dprint(fg)
pytensor.config.optimizer_verbose = True
for i in range(3):    
    print()
    print(">>> ", end="")
    pytensor.dprint(history.next())
        
# Log [id A] 4
#  └─ True_div [id B] 3
#     ├─ Exp [id C] 2
#     │  └─ x [id D]
#     └─ Sum{axes=None} [id E] 1
#        └─ Exp [id F] 0
#           └─ x [id D]
# >>> MergeOptimizer
# Log [id A] 3
#  └─ True_div [id B] 2
#     ├─ Exp [id C] 0
#     │  └─ x [id D]
#     └─ Sum{axes=None} [id E] 1
#        └─ Exp [id C] 0
#           └─ ···
# >>> local_mul_canonizer
# Log [id A] 1
#  └─ Softmax{axis=None} [id B] 0
#     └─ x [id C]
# >>> local_logsoftmax
# LogSoftmax{axis=None} [id A] 0
#  └─ x [id B]

# Or in reverse
for i in range(3):
    print()
    print(">>> ", end="")
    pytensor.dprint(history.prev())

# >>> local_logsoftmax
# Log [id A] 1
#  └─ Softmax{axis=None} [id B] 0
#     └─ x [id C]
# >>> local_mul_canonizer
# Log [id A] 3
#  └─ True_div [id B] 2
#     ├─ Exp [id C] 0
#     │  └─ x [id D]
#     └─ Sum{axes=None} [id E] 1
#        └─ Exp [id C] 0
#           └─ ···
# >>> MergeOptimizer
# Log [id A] 4
#  └─ True_div [id B] 3
#     ├─ Exp [id C] 2
#     │  └─ x [id D]
#     └─ Sum{axes=None} [id E] 1
#        └─ Exp [id F] 0
#           └─ x [id D]
    
pytensor.config.optimizer_verbose = False
# Or go to any step
pytensor.dprint(history.goto(2))
# Log [id A] 1
#  └─ Softmax{axis=None} [id B] 0
#     └─ x [id C]

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

Suggested change

	from pytensor.graph.features import History, FullHistory
	from pytensor.graph.features import FullHistory

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB