Any way to workaround "No replication rule"?

[southfreebird]

Hi everyone!

I have a replicated Pallas/Triton kernel similar to identity. I'm sharding it with shard_map using check_rep=False, but this leads to unwanted AllReduce operations in the background. Is there something I can do about this?

Here are my thoughts:

• Implement a replication rule for my custom_call to make it possible to use shard_map with check_rep=True. Is this even possible?
• Implement custom partitioning for my kernel? Will this help in my case?
• Any other ideas?