Explicitly raise an error if more than 65535 channels are created

`xla::HostCallbackArgInfo` uses `uint16_t` for channel ids, so we should warn users explicitly when the channel ids exceed the UINT16_MAX instead of silently wrapping around.

PiperOrigin-RevId: 695682871

Author: junwhanahn

jax/_src/interpreters/mlir.py

@@ -770,7 +770,13 @@ def backend(self) -> xb.XlaBackend:
     return self.backend_or_name
 
   def new_channel(self) -> int:
-    return next(self.channel_iterator)
+    channel = next(self.channel_iterator)
+    # `xla::HostCallback` requires a 16-bit channel ID.
+    if channel >= (1 << 16):
+      raise RuntimeError(
+          "Host callback lowering created too many channels. PjRt does not"
+          " support more than 65535 channels")
+    return channel
 
   # Adds an IFRT host callback object to the context. A reference to these
   # callbacks will be provided to IFRT during compilation so it can do things