Correct bad imports in optimize.py and expose it via pytensor.tensor.__init__

[jessegrabowski]

Description

Small PR to improve the visibility of pytensor.tensor.optimize and to fix some bad imports in the file.

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

---

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

[ricardoV94]

Also add it to the docs?

[jessegrabowski]

Yes you are right of course.

[codecov]

Codecov Report

Attention: Patch coverage is 64.86486% with 13 lines in your changes missing coverage. Please review.

Project coverage is 82.01%. Comparing base (646a734) to head (e073022).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/optimize.py	64.86%	10 Missing and 3 partials ⚠️

❌ Your patch status has failed because the patch coverage (64.86%) 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    #1464      +/-   ##
==========================================
- Coverage   82.14%   82.01%   -0.13%     
==========================================
  Files         213      214       +1     
  Lines       50270    50416     +146     
  Branches     8879     8901      +22     
==========================================
+ Hits        41295    41350      +55     
- Misses       6774     6858      +84     
- Partials     2201     2208       +7     

Files with missing lines	Coverage Δ
pytensor/tensor/optimize.py	87.30% <64.86%> (-3.52%)	⬇️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[ricardoV94]

Let's benchmark the import time of PyTensor to see if importing scipy optimize doesn't cause a big slowdown

[ricardoV94]

python -m benchmark_imports pytensor is okay enough we should check with and without the optimize import in init

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[ricardoV94]

Does it look nicer with jacobian vectorize?

Reply via ReviewNB

[jessegrabowski]

Is that merged? Or you want me to try the nb in that feature branch

[jessegrabowski]

python -m benchmark_imports pytensor is okay enough we should check with and without the optimize import in init

I don't understand what this means

[ricardoV94]

python -m benchmark_imports pytensor

Install https://pypi.org/project/benchmark-imports and check the output of that line before and after you made optimize.py imported in __init__. I'm worried that scipy optimize import can be a non-neglible overhead

[ricardoV94]

This was the case with the eager import of scipy.stats: #1268

[jessegrabowski]

Here is the output:

0.5977 root       pytensor                                
0.2493 project    pytensor.configdefaults                 
0.1462 project    pytensor.scalar                         
0.1238 project    pytensor.scalar.math                    
0.1145 dependency scipy.special                           
0.1123 project    pytensor.tensor                         
0.0710 project    pytensor.tensor.optimize                
0.0684 dependency scipy.optimize                          
0.0526 transitive scipy.special._orthogonal                from scipy.special
0.0522 transitive scipy.linalg                             from scipy.special._orthogonal
0.0487 transitive scipy.special._support_alternative_backends from scipy.special
0.0484 transitive scipy._lib._array_api                    from scipy.special._support_alternative_backends
0.0466 transitive scipy._lib.array_api_compat.numpy        from scipy._lib._array_api
0.0353 project    pytensor.tensor.rewriting               
0.0348 transitive scipy.linalg._sketches                   from scipy.linalg
0.0343 transitive scipy.sparse                             from scipy.linalg._sketches
0.0292 dependency numpy                                   
0.0243 transitive scipy.sparse.csgraph                     from scipy.sparse
0.0221 project    pytensor.scalar.basic                   
0.0201 transitive scipy.optimize._shgo                     from scipy.optimize
0.0192 project    pytensor.printing                       
0.0184 project    pytensor.compile                        
0.0178 project    pytensor.compile.function               
0.0177 transitive numpy.f2py                               from scipy._lib.array_api_compat.numpy
0.0175 transitive scipy.spatial                            from scipy.optimize._shgo

I'm happy to move the imports into perform though, even if the cost is small

[ricardoV94]

Can you show the output with the lazy optimize (just revert adding it to __init__.py to test

[ricardoV94]

Cost seems to be around 1.2x slower. Also keeping in mind the majority of users will not use these Ops or may use only in other backends

[jessegrabowski]

True enough. I moved imports into the perform methods. Here's the updated benchmark:

0.7534 root       pytensor                                
0.4372 project    pytensor.configdefaults                 
0.1705 project    pytensor.scalar                         
0.1461 project    pytensor.scalar.math                    
0.1318 dependency scipy.special                           
0.0588 transitive scipy.special._orthogonal                from scipy.special
0.0580 transitive scipy.linalg                             from scipy.special._orthogonal
0.0577 project    pytensor.tensor                         
0.0568 transitive scipy.special._support_alternative_backends from scipy.special
0.0564 transitive scipy._lib._array_api                    from scipy.special._support_alternative_backends
0.0540 transitive scipy._lib.array_api_compat.numpy        from scipy._lib._array_api
0.0456 dependency numpy                                   
0.0440 project    pytensor.tensor.rewriting               
0.0377 transitive scipy.linalg._sketches                   from scipy.linalg
0.0371 transitive scipy.sparse                             from scipy.linalg._sketches
0.0254 transitive scipy.sparse.csgraph                     from scipy.sparse
0.0240 project    pytensor.scalar.basic                   
0.0220 transitive numpy.__config__                         from numpy
0.0215 transitive numpy._core                              from numpy.__config__
0.0210 project    pytensor.printing                       
0.0202 project    pytensor.compile                        
0.0198 transitive numpy.f2py                               from scipy._lib.array_api_compat.numpy
0.0196 project    pytensor.compile.function               
0.0194 transitive numpy.testing                            from scipy._lib.array_api_compat.numpy
0.0193 transitive numpy.f2py.f2py2e                        from numpy.f2py

[ricardoV94]

We should also remove the props. you can override the str/repr methods to show them, but we shouldn't use props because they are not enough to ascertain equality of ops (the inner graph would also have to be considered in the hash and equality).

This is something we still need to work on