From 2dba9cfed7a330f3e68e907ade836e97a36a3f41 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 3 Oct 2025 09:01:11 -0700 Subject: [PATCH 1/2] xla_home logging --- build_tools/jax.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index 1f9552eb69..c3f5558642 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -36,15 +36,17 @@ def xla_path() -> str: from jax.extend import ffi # pylint: disable=ungrouped-imports except ImportError: + print("Could not import jax. Looking for XLA source in $XLA_HOME or /opt/xla") if os.getenv("XLA_HOME"): xla_home = Path(os.getenv("XLA_HOME")) else: xla_home = "/opt/xla" else: xla_home = ffi.include_dir() + print(f"Found XLA source in {xla_home}") if not os.path.isdir(xla_home): - raise FileNotFoundError("Could not find xla source.") + raise FileNotFoundError(f"Could not find xla source. Searched: {xla_home}") return xla_home From b455f7a340b0c24176918c625845b79087ea301a Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 3 Oct 2025 09:17:56 -0700 Subject: [PATCH 2/2] fix --- build_tools/jax.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index c3f5558642..b03c7e59d3 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -38,7 +38,7 @@ def xla_path() -> str: except ImportError: print("Could not import jax. Looking for XLA source in $XLA_HOME or /opt/xla") if os.getenv("XLA_HOME"): - xla_home = Path(os.getenv("XLA_HOME")) + xla_home = Path(os.getenv("XLA_HOME").strip()) else: xla_home = "/opt/xla" else: