Where is TPU device classification defined in JAX?

[samixyzdev]

Where is TPU device classification defined in JAX?

Hi JAX team and contributors,

I'm looking to contribute to JAX by modifying how TPU devices are classified.
Currently, multi-slice TPU devices are labeled as MegaScalePjRtDevice, but I believe they should remain TpuDevice.

I've tried tracing the classification process and followed these steps:

• Searched for MegaScalePjRtDevice in the JAX codebase.
• Traced TPU backend creation from make_tpu_client(), which is called in xla_bridge.py.
• Found that make_tpu_client() is inside jaxlib, likely in a compiled .so file.
• Looked for libtpu.so but didn’t find it on my system.
• Checked if TPU support was loaded dynamically via pjrt_plugin, but I couldn't confirm its exact behavior.

Question: Where exactly in the JAX codebase is the TPU device classification (MegaScalePjRtDevice vs. TpuDevice) handled? I'd love to contribute a fix, but I need to understand where this happens.

Any guidance is appreciated! Thanks in advance. 😊