Set XLA_EXTENSION_DIR on build for xla-rs to read

Cargo.toml

@@ -22,7 +22,7 @@ slotmap = "1"
 backtrace = "0.3"
 smallvec = "1.13"
 strum_macros = "0.26"
-xla = { git = "https://github.com/unda-ai/xla-rs", version = "0.1.6" , branch = "main" }
+xla = { git = "https://github.com/unda-ml/xla-rs", version = "0.1.6" , branch = "main" }
 thiserror = "1"
 half = "2.4.0"
 byteorder = "1.5"

build.rs

@@ -0,0 +1,13 @@
+use std::env;
+use std::path::PathBuf;
+
+fn main() {
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").expect("Missing manifest");
+    let xla_extension_dir = PathBuf::from(manifest_dir).join("xla_extension");
+
+    if !xla_extension_dir.exists() {
+        panic!("The xla_extension library was not found in the expected path: {}", xla_extension_dir.display());
+    }
+    println!("cargo:rustc-env=XLA_EXTENSION_DIR={}", xla_extension_dir.display());
+    println!("cargo:rerun-if-changed=build.rs");
+}