Skip to content

Commit e7eb207

Browse files
hawkinspjax authors
authored andcommitted
Change xla_bridge_test to expect a bytes FDO profile instead of a string in future jaxlib versions.
Change in preparation for nanobind migration. PiperOrigin-RevId: 613329489
1 parent ef8f4fc commit e7eb207

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

tests/xla_bridge_test.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from jax._src import xla_bridge as xb
2727
from jax._src.interpreters import xla
2828
from jax._src.lib import xla_client as xc
29+
from jax._src.lib import xla_extension_version
2930

3031
config.parse_flags_with_absl()
3132

@@ -55,7 +56,8 @@ def test_set_fdo_profile(self):
5556
num_replicas=1, num_partitions=1, fdo_profile=b"test_profile"
5657
)
5758
self.assertEqual(
58-
compile_options.executable_build_options.fdo_profile, "test_profile"
59+
compile_options.executable_build_options.fdo_profile,
60+
b"test_profile" if xla_extension_version >= 242 else "test_profile"
5961
)
6062

6163
def test_autofdo_profile(self):
@@ -386,6 +388,5 @@ def test_experimental_warning(self):
386388
)
387389

388390

389-
390391
if __name__ == "__main__":
391392
absltest.main(testLoader=jtu.JaxTestLoader())

0 commit comments

Comments
 (0)