Merge pull request #16392 from cloudhan:fix-multi-line

PiperOrigin-RevId: 540691816
google · Jun 15, 2023 · 904b46a · 904b46a
2 parents 57e5d08 + 7b8bdfb
commit 904b46a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/jax/_src/lax/control_flow/common.py b/jax/_src/lax/control_flow/common.py
@@ -41,6 +41,9 @@ def _typecheck_param(prim, param, name, msg_required, pred):
     msg = (f'invalid {prim} param {name} of type {type(param).__name__}, '
            f'{msg_required} required:')
     param_str = str(param)
+    # Avoid using os.linesep here to have the same multi-line error message
+    # format on different platforms.
+    sep = '\n' if '\n' in param_str or '\r' in param_str else ' '
     sep = os.linesep if os.linesep in param_str else ' '
     msg = sep.join([msg, param_str])
     raise core.JaxprTypeError(msg)