What happened to the XPU backend?

[dionhaefner]

In mpi4jax, we noticed that it's not possible to register XPU extensions anymore. Was that backend cut from JAX?

NotImplementedError: Registering an MLIR lowering rule for primitive sendrecv_mpi for an unknown platform xpu. Known platforms are: METAL, cpu, cuda, gpu, rocm, tpu.

[hawkinsp]

We've never had an in-tree XPU backend.

We added checking code that means you can't register lowering rules for backends we've never heard of. If you have the XPU plugin installed, then we should allow such rules to be registered. You could also add xpu here

jax/jax/_src/xla_bridge.py

Line 206 in 6235158

_experimental_plugins: set[str] = {"METAL"}

so we allow it even if there's no such plugin installed.

The goal of the check is mostly to catch typos, not to forbid new backends.

[dionhaefner]

That makes sense, thanks