Unrecognisable XLA flags in JAX built on TPU with custom XLA

[awav]

Hello everyone, this question relates to the #9220 (thanks to @skye and @zhangqiaorjc for resolving it quickly). Now, I'm puzzled how to make JAX to run my custom built XLA with extra optimisation passes on a TPU VM. For CPU and GPU it was straightforward - I simply added my optimisation passes in (cpu|gpu)_compiler.cc. I suspect I need to do something similar with TPU backend, but I can't figure out what exactly.

I would really appreciate it if someone could help me with resolving this issue. Thanks!

PS: Error I'm getting right now (with dummy custom flags):

2022-01-20 16:55:41.526525: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:171] ****(awav) THIS IS A CUSTOM BUILD****
2022-01-20 16:55:41.526573: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:172] XLA service 0x1b83380 initialized for platform Interpreter (this does not guarantee that XLA will be used). Devices:
2022-01-20 16:55:41.526582: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:180]   StreamExecutor device (0): Interpreter, <undefined>
2022-01-20 16:55:41.539829: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.cc:164] TfrtCpuClient created.
INFO:absl:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker:
INFO:absl:Unable to initialize backend 'gpu': NOT_FOUND: Could not find registered platform with name: "cuda". Available platform names are: TPU Host Interpreter
F0120 16:55:41.547101  190667 parse_flags_from_env.cc:229] Unknown flags in XLA_FLAGS: --xla_custom_flag=hello_world Perhaps you meant to specify these on the TF_XLA_FLAGS envvar?

[zhangqiaorjc]

Sorry for the delay.

F0120 16:55:41.547101  190667 parse_flags_from_env.cc:229] Unknown flags in XLA_FLAGS: --xla_custom_flag=hello_world Perhaps you meant to specify these on the TF_XLA_FLAGS envvar?

Did you try TF_XLA_FLAGS?

TF_XLA_FLAGS=--xla_custom_flag =hell_world path/to/your/program

[ayaka14732]

Not sure but there is also LIBTPU_INIT_ARGS=--xla_custom_flag=hello_world