Request description
I found the verification logic for composite ops here, and it looks correct to me.
However, I’m wondering what your thoughts are on applying sharding to composite ops.
As far as I can tell, there’s currently no defined rule for sharding composite ops in the Shardy repo.
Let’s assume we can apply sharding in a pointwise manner — what would a sharded composite function look like in that case, particularly in terms of its function signature?
I experimented by forcing a composite function to be sharded, but encountered type mismatch errors due to inconsistencies between the function signature and the sharded input/output tensors.
I can manually modify the function signature to match the sharded tensors, but I’m not sure if that’s the correct approach.
Ideally, I was hoping that StableHLO’s verification could recognize sharding annotations, allowing the verification to pass more naturally.
If you have any ideas on this, I’d really appreciate your input.
Also, if there’s any ongoing or planned work to support sharding on composite ops, I’d love to hear about it.
Request description
I found the verification logic for composite ops here, and it looks correct to me.
However, I’m wondering what your thoughts are on applying sharding to composite ops.
As far as I can tell, there’s currently no defined rule for sharding composite ops in the Shardy repo.
Let’s assume we can apply sharding in a pointwise manner — what would a sharded composite function look like in that case, particularly in terms of its function signature?
I experimented by forcing a composite function to be sharded, but encountered type mismatch errors due to inconsistencies between the function signature and the sharded input/output tensors.
I can manually modify the function signature to match the sharded tensors, but I’m not sure if that’s the correct approach.
Ideally, I was hoping that StableHLO’s verification could recognize sharding annotations, allowing the verification to pass more naturally.
If you have any ideas on this, I’d really appreciate your input.
Also, if there’s any ongoing or planned work to support sharding on composite ops, I’d love to hear about it.