From 03876bd702140197d0bc1a1ea48af75f29fd266f Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 16 Aug 2022 15:47:15 +0000
Subject: [PATCH] build.py fixes.

* Add aarch64 as a known target_cpu value.
* Only pass --bazel_options to build actions since they can make "bazel
  shutdown" fail.
* Pass the bazel startup options to "bazel shutdown".

Issue https://github.com/google/jax/issues/7097
Fixes https://github.com/google/jax/issues/7639
---
 build/build.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/build/build.py b/build/build.py
index 96fb3d27777b..c37d57153242 100755
--- a/build/build.py
+++ b/build/build.py
@@ -266,7 +266,7 @@ def write_bazelrc(*, python_bin_path, remote_build,
       f.write("build --distinct_host_configuration=false\n")
 
     for o in bazel_options:
-      f.write(f"common {o}\n")
+      f.write(f"build {o}\n")
     if target_cpu_features == "release":
       if wheel_cpu == "x86_64":
         f.write("build --config=avx_windows\n" if is_windows()
@@ -464,6 +464,7 @@ def main():
       "darwin_arm64": "arm64",
       "darwin_x86_64": "x86_64",
       "ppc": "ppc64le",
+      "aarch64": "aarch64",
   }
   # TODO(phawkins): support other bazel cpu overrides.
   wheel_cpu = (wheel_cpus[args.target_cpu] if args.target_cpu is not None
@@ -548,7 +549,7 @@ def main():
     f"--cpu={wheel_cpu}"])
   print(" ".join(command))
   shell(command)
-  shell([bazel_path, "shutdown"])
+  shell([bazel_path] + args.bazel_startup_options + ["shutdown"])
 
 
 if __name__ == "__main__":