Prevent the XLA compiler from sharding the custom call in favour of Mosaic sharding based on user annotations.

PiperOrigin-RevId: 614336455

Author: naummo

jax/_src/tpu_custom_call.py

@@ -141,6 +141,9 @@ def to_json(self) -> bytes:
         if i + 1 != len(self.flags):
           config.write(b",")
       config.write(b"]")
+    # Prevent the compiler from sharding the custom call beyond what Mosaic does
+    # based on user annotations
+    config.write(b', "implicit_sharding": {"type": "MANUAL"}')
     config.write(b"}")
     return config.getvalue()