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. |
| if shaped_inputs[0].dim() == 1: | ||
| ufunc = torch.vmap(base_fn) | ||
| else: | ||
| dims = (tuple(range(shaped_inputs[0].dim())),) |
There was a problem hiding this comment.
I don't think this is correct. You need to apply vmap repeatedly to make it vectorize across multiple dimensions
There was a problem hiding this comment.
You're correct, my bad. I misunderstood the docs and wrote a bad local test.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1032 +/- ##
==========================================
- Coverage 81.90% 81.89% -0.01%
==========================================
Files 182 182
Lines 47879 47887 +8
Branches 8620 8619 -1
==========================================
+ Hits 39214 39216 +2
- Misses 6492 6498 +6
Partials 2173 2173
|
| # @todo: This will fail for anything that calls | ||
| # `.item()` |
There was a problem hiding this comment.
We usually use # TODO: not # @todo:
| return base_fn(*inputs) | ||
| else: | ||
|
|
||
| def elemwise_fn(*inputs): |
There was a problem hiding this comment.
We need a test for this branch. You can create a new test ScalarOp and dispatch something specifically in PyTorch (in the test suite) to test this?
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/