Refactor `infer_shape` method of Ops to find output shapes using `gufunc_signature`

Open Aarsh-Wankar opened this issue 9 months ago • 3 comments

Refactored infer_shape method of Ops to find output shapes using gufunc_signature using a newly defined function _gufunc_to_out_shape.

Description

Ops have a method infer_shape which helps to find the shapes of the outputs given the input shapes. All Ops have their own implementations of infer_shape. However, many Ops have a gufunc_signature string which gives information about the input-output shape relations. In principle, this string is enough to find the output shapes given the input shapes. Thus, a function _gufunc_to_out_shape has been added in this PR, which calculates the output shapes list given the gufunc_signature and the input shapes list. This PR also replaces the Op specific implementations of infer_shape with the _gufunc_to_out_shape output.

Related Issue

[x] Closes #1257
[ ] Related to #

Checklist

[x] Checked that the pre-commit linting/style checks pass
[ ] Included tests that prove the fix is effective or that the new feature works
[x] 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
[x] Maintenance
[ ] Other (please specify):

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

Mar 13 '25 15:03 Aarsh-Wankar