Allow for more elemwise torch functions using broadcast_tensor and vmap
#1032
Conversation
|
I need to add a test, but I want to get feedback on #1031 before continuing. |
Ch0ronomato
changed the title
broadcast_tensor and vmap
|
I'll fix the tests. |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1032 +/- ##
==========================================
- Coverage 82.10% 82.10% -0.01%
==========================================
Files 183 183
Lines 47924 47932 +8
Branches 8632 8634 +2
==========================================
+ Hits 39348 39354 +6
- Misses 6410 6413 +3
+ Partials 2166 2165 -1
|
|
|
||
| def elemwise_fn(*inputs): | ||
| Elemwise._check_runtime_broadcast(node, inputs) | ||
| shaped_inputs = torch.broadcast_tensors(*inputs) |
There was a problem hiding this comment.
Nit: more precise name:
| shaped_inputs = torch.broadcast_tensors(*inputs) | |
| broadcasted_inputs = torch.broadcast_tensors(*inputs) |
Also needs to be changed below
| # @todo: This will fail for anything that calls | ||
| # `.item()` |
There was a problem hiding this comment.
Remove todo, not something that should be addressed in this impl, but for specific Ops, so if anything should exist as a github issue?
There was a problem hiding this comment.
nope i agree with you.
|
I loosen the failing test in #988 ; this should be good after that. |
| def elemwise_fn(*inputs): | ||
| Elemwise._check_runtime_broadcast(node, inputs) | ||
| return base_fn(*inputs) | ||
| if hasattr(scalar_op, "nfunc_spec") and hasattr(torch, scalar_op.nfunc_spec[0]): |
There was a problem hiding this comment.
Should do the same trick for scipy.x you did in another PR?
There was a problem hiding this comment.
...but then I have to write a new test for this PR 😅
Ch0ronomato
force-pushed
the
elemwise_torch_improvement
branch
2 times, most recently
from
97d6bdc to
1dd6322
Compare
Ch0ronomato
force-pushed
the
elemwise_torch_improvement
branch
from
1dd6322 to
c7e8a64
Compare
Description
In the event the operator
Elemwiseis broadcasting over doesn't have a direct torch function, we can leveragevmapandbroadcast_tensorsto replicate the ufunc machinery.Related Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1032.org.readthedocs.build/en/1032/