Add print_shape and print_memory_map option to debugprint

[Aarsh-Wankar]

Description

I added an extra argument print_shape to debugprint and _debugprint, which, when True, prints the shape of every Variable in the graph.

Test demonstration:

import pytensor
import pytensor.tensor as pt
x = pt.matrix("x", shape=(None, 3))
y = pt.broadcast_to(x * 5, (5, 7, 3))
y.dprint(print_shape = True)

# Output:
# Alloc [id A] shape=(5, 7, 3)
#  ├─ Mul [id B] shape=(?, 3)
#  │  ├─ x [id C] shape=(?, 3)
#  │  └─ ExpandDims{axes=[0, 1]} [id D] shape=(1, 1)
#  │     └─ 5 [id E] shape=()
#  ├─ 5 [id F] shape=()
#  ├─ 7 [id G] shape=()
#  └─ 3 [id H] shape=()

Additionally, I also simplified the __str__ function logic of class TensorType to easily find the formatted shape.

formatted_shape = str(shape).replace("None", "?")

Related Issue

• Closes Add a print_shape argument to debug_print #1191
• 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):

---

📚 Documentation preview 📚: https://pytensor--1236.org.readthedocs.build/en/1236/

[codecov]

Codecov Report

Attention: Patch coverage is 50.00000% with 5 lines in your changes missing coverage. Please review.

Project coverage is 81.98%. Comparing base (5d4e9e0) to head (640c759).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/printing.py	37.50%	3 Missing and 2 partials ⚠️

❌ Your patch status has failed because the patch coverage (50.00%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1236      +/-   ##
==========================================
- Coverage   81.99%   81.98%   -0.02%     
==========================================
  Files         188      188              
  Lines       48566    48567       +1     
  Branches     8677     8677              
==========================================
- Hits        39823    39818       -5     
- Misses       6580     6584       +4     
- Partials     2163     2165       +2     

Files with missing lines	Coverage Δ
pytensor/tensor/type.py	94.52% <100.00%> (-0.08%)	⬇️
pytensor/printing.py	51.43% <37.50%> (-0.25%)	⬇️

... and 1 file with indirect coverage changes

[ricardoV94]

Awesome!

[ricardoV94]

Another thing that I've been missing is to have a print_memory_map which is equivalent to setting both print_destroy_map and print_view_map to True. Usually I need both or neither, never just the one

[ricardoV94]

Thanks @Aarsh-Wankar