how to compile jax with custom devices

[HardysJin]

Hi,

Hope u are doing well. I have a modified xla with support of custom NPU devices. I am wondering how to add the custom backend and let jax recognize it by e.g. jax.devices('npu'). Of course it should compile with

python build/build.py --bazel_options=--override_repository=xla=/path/to/xla-npu

but jax.devices('npu') returns RuntimeError: Unknown backend npu. Available backends are ['cpu']

I would really appreciate it if someone could help me with resolving this issue. Thank you so much!

Best,
Hardys

[thoth291]

Was it ever resolved in any way?