File tree Expand file tree Collapse file tree 1 file changed +3
-2
lines changed Expand file tree Collapse file tree 1 file changed +3
-2
lines changed Original file line number Diff line number Diff line change 26
26
from jax ._src import xla_bridge as xb
27
27
from jax ._src .interpreters import xla
28
28
from jax ._src .lib import xla_client as xc
29
+ from jax ._src .lib import xla_extension_version
29
30
30
31
config .parse_flags_with_absl ()
31
32
@@ -55,7 +56,8 @@ def test_set_fdo_profile(self):
55
56
num_replicas = 1 , num_partitions = 1 , fdo_profile = b"test_profile"
56
57
)
57
58
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"
59
61
)
60
62
61
63
def test_autofdo_profile (self ):
@@ -386,6 +388,5 @@ def test_experimental_warning(self):
386
388
)
387
389
388
390
389
-
390
391
if __name__ == "__main__" :
391
392
absltest .main (testLoader = jtu .JaxTestLoader ())
You can’t perform that action at this time.
0 commit comments