Generic batch op interface #151
Conversation
|
I'm dubious if we should go this way? AFAIK Jax doesn't do this unrolling. Like batching over a lot of operations should still be done without unrolling. Do you know for which ops should the unrolling be done? |
|
Another possible generic strategy (maybe as a fallback) is to wrap the op in a |
|
Here I have enabled the unrolling for just a few ops where the batch interface is needed and not implemented. |
src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
Outdated
Show resolved
Hide resolved
Pangoraw
changed the title
| auto iszero = matchPattern(ifOp.getPred(), m_Zero()); | ||
| auto isone = matchPattern(ifOp.getPred(), m_One()); | ||
|
|
||
| if (!iszero && !isone) |
There was a problem hiding this comment.
can we also test if the input operand would be a broadcastop of a single value?
There was a problem hiding this comment.
added it but I am having a hard time coming up with an input to test this where the batched version of the predicate is a broadcast_in_dim.
There was a problem hiding this comment.
yeah I think long term this will require an extension to the batching interface to mark some operands as non-batched, in any case is fine for now
3 participants
Depends on EnzymeAD/Enzyme#2142.