use "true" and "false" instead of 0 and 1 #890
Conversation
Jax/XLA seem to expect true and false strings instead of using 0 and 1.
| @@ -24,10 +24,10 @@ def f(x: Tensor) -> Tensor: | |||
| ) | |||
| self.assertEqual(f_compiled(5), 15) | |||
|
|
|||
| def atest_xla_flags_from_options(self): | |||
There was a problem hiding this comment.
this test was never being run because of the a before test in function name.
|
I updated the existing test and renamed the test so it's being run correctly now. |
| @@ -144,7 +144,7 @@ def xla_flags_from_options(xla_options: dict[str, Union[str, bool, int]]) -> str | |||
| flags = [] | |||
| for k, v in xla_options.items(): | |||
| if isinstance(v, bool): | |||
| v = "1" if v else "0" | |||
| v = "true" if v else "false" | |||
There was a problem hiding this comment.
Should we convert all the "true" and "false" strings in default_xla_options to bools to be consistent?
There was a problem hiding this comment.
I think that should be preferred. I just wanted to keep this PR small to make merging and reviewing easier. Happy to include that change in this PR or do a follow up.
There was a problem hiding this comment.
Thanks, I think it makes sense to include in this PR.
| @@ -144,7 +144,7 @@ def xla_flags_from_options(xla_options: dict[str, Union[str, bool, int]]) -> str | |||
| flags = [] | |||
| for k, v in xla_options.items(): | |||
| if isinstance(v, bool): | |||
| v = "1" if v else "0" | |||
| v = "true" if v else "false" | |||
There was a problem hiding this comment.
Would it make sense to do an aot compilation or run the trainer for a few steps to confirm these changes work? (If you haven't already?)
There was a problem hiding this comment.
I did test it on v6e afair last year, but any additional runs would be great.
There was a problem hiding this comment.
Maybe test it again on v5p using the jax/jaxlib version that is currently pinned in AXLearn? (0.4.33 IIRC)
Jax/XLA expects true and false strings instead of using 0 and 1 for certain kind of boolean flags (TriStateFlag which also accepts "auto").
This should solve this issue:
Confirmed by Jax team that we should use "true" and "false" so TriStateflags which are like boolean flags work as well.