Implement C code for ExtractDiagonal and ARange #1392
Conversation
Set view flag of ExtractDiagonal to True and respect by default
Codecov ReportAttention: Patch coverage is
❌ Your patch check has failed because the patch coverage (89.18%) 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 #1392 +/- ##
==========================================
+ Coverage 82.02% 82.04% +0.01%
==========================================
Files 207 207
Lines 49294 49305 +11
Branches 8746 8745 -1
==========================================
+ Hits 40433 40450 +17
+ Misses 6695 6689 -6
Partials 2166 2166
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull Request Overview
This PR implements C code for two tensor operations: ARange and ExtractDiag, addressing numerical consistency and performance. Key changes include subclassing both operations from COp (instead of Op) for C-level performance improvements, updating ARange’s perform and c_code methods, and modifying ExtractDiag to enable non‐writeable views by default with safe fallbacks.
| """Create an array containing evenly spaced values within a given interval. | ||
|
|
||
| Parameters and behaviour are the same as numpy.arange(). | ||
|
|
||
| """ | ||
|
|
||
| # TODO: Arange should work with scalars as inputs, not arrays |
There was a problem hiding this comment.
[nitpick] Consider addressing the TODO by supporting scalar inputs for ARange, which would simplify the interface and align behavior with numpy.arange.
Copilot uses AI. Check for mistakes.
| try: | ||
| out.flags.writeable = True | ||
| except ValueError: | ||
| # We can't make this array writable |
There was a problem hiding this comment.
[nitpick] Consider logging a warning or providing additional context here when the array cannot be made writeable and a copy is made, to aid in diagnosing potential performance implications.
| # We can't make this array writable | |
| # We can't make this array writable | |
| warnings.warn( | |
| "Unable to make the array writeable. A copy of the array is being created instead. " | |
| "This may have performance implications if the array is large.", | |
| RuntimeWarning, | |
| ) |
Copilot uses AI. Check for mistakes.
Closes #1360.
It also allows ExtractDiagonal to be a view (it's on by default). Under the hood numpy has returned non-writeable views since 1.9.0, and we have a lower pin of 1.17.0. It's as simple as swapping the flag [citation needed].
📚 Documentation preview 📚: https://pytensor--1392.org.readthedocs.build/en/1392/