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?

Screencast.From.2025-06-10.19-03-08.mp4

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):

[codecov]

Codecov Report

❌ Patch coverage is 35.93750% with 82 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.03%. Comparing base (d10f245) to head (84e56ce).
⚠️ Report is 136 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/ipython.py	0.00%	78 Missing ⚠️
pytensor/graph/features.py	92.00%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #874      +/-   ##
==========================================
- Coverage   82.12%   82.03%   -0.09%     
==========================================
  Files         211      214       +3     
  Lines       49757    50398     +641     
  Branches     8819     8897      +78     
==========================================
+ Hits        40862    41345     +483     
- Misses       6715     6848     +133     
- Partials     2180     2205      +25     

Files with missing lines	Coverage Δ
pytensor/graph/features.py	67.63% <92.00%> (+3.06%)	⬆️
pytensor/ipython.py	0.00% <0.00%> (ø)

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull Request Overview

Adds a new FullHistory feature to record and navigate rewrite steps on a FunctionGraph, along with a corresponding test and an IPython widget for interactive exploration.

• Implements FullHistory in pytensor/graph/features.py to capture and rewind/fast-forward graph rewrites
• Adds test_full_history in tests/graph/test_features.py and updates imports for rewrite testing
• Introduces InteractiveRewrite widget in pytensor/ipython.py, and adjusts test/CI configs to skip it

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
pytensor/graph/features.py	New FullHistory feature for tracking and replaying rewrites
tests/graph/test_features.py	Updated imports and added test_full_history
pytensor/ipython.py	New interactive widget (InteractiveRewrite) for notebook use
pyproject.toml	Excludes pytensor/ipython.py from pytest runs
.github/workflows/test.yml	Updates CI to ignore pytensor/ipython.py in doctests

Comments suppressed due to low confidence (4)

pytensor/graph/features.py:536

• [nitpick] The internal list name 'fw' is ambiguous; consider renaming it to something more descriptive like 'forward_steps' or 'forward_history' for better readability.

self.fw = []

pytensor/graph/features.py:554

• [nitpick] Add a detailed docstring to goto explaining the valid range of checkpoint, how it maps to history states, and its effect on pointer.

def goto(self, checkpoint):

pytensor/ipython.py:1

• The new InteractiveRewrite widget isn’t covered by any tests; consider adding a basic smoke test (e.g., instantiating the widget and verifying core attributes) to prevent regressions.

import anywidget

pytensor/ipython.py:6

• Verify that Variable (and rewrite_graph) are exported from pytensor.graph; if not, import Variable from pytensor.graph.basic and rewrite_graph from pytensor.graph.rewriting.utils to avoid import errors.

from pytensor.graph import FunctionGraph, Variable, rewrite_graph

[Copilot]

Using assert for input validation can be bypassed in optimized runs; consider raising a ValueError for out-of-range checkpoint values instead.

Suggested change

	assert 0 <= checkpoint <= history_len
	if not (0 <= checkpoint <= history_len):
	raise ValueError(f"Checkpoint value {checkpoint} is out of range. It must be between 0 and {history_len}.")

[zaxtax]

LGTM overall. It does feel like it should have a few more unit tests, but unsure concretely which ones to add.

[zaxtax]

I think these should be put in a widget.js and a widget.css file instead of hanging out as a giant string.

[ricardoV94]

For now I would keep it together, until and if it actually grows larger. It's a single file this way instead of a folder. We can always refactor later