jax
jax copied to clipboard
jaxlib should not link in tpu_initializer_helper when TPU is disabled
The tpu_initializer_helper should not be linked in if TPU support is disabled:
$ python build/build.py --configure_only
...
TPU enabled: no
Remote TPU enabled: no
...
$ bazel cquery 'somepath(//build:build_wheel,@org_tensorflow//tensorflow/core/tpu:tpu_initializer_helper)'
....
//build:build_wheel (e92e97f)
//jaxlib:jaxlib (e92e97f)
//jaxlib:xla_extension (e92e97f)
@org_tensorflow//tensorflow/compiler/xla/python:xla_extension.so (e92e97f)
@org_tensorflow//tensorflow/compiler/xla/pjrt:pjrt_c_api_client (e92e97f)
@org_tensorflow//tensorflow/core/tpu:tpu_initializer_helper (e92e97f)
...
This is breaking Windows builds but it is also superfluous on other platforms that don't have TPUs, e.g., Mac.
See #12443