Add blacken script and format fixes

blacken.sh can be used by developers to run
format on the python source files in the repo.

The scope of files it runs on is limited at the moment
since running black on the jax files results in large
diffs which makes it hard to track changes from upstream.
Eventually when the PJRT stuff is fully integrated to the repo
we can run it on those. The jaxlib kernels is still an
open question.

The script will pass any args it receives to the inner black
command so args like --diff and --check work.

Author: mrodden

build/ci_build

@@ -153,9 +153,7 @@ def test(image_name, test_cmd=None):
     cmd.extend(mounts)
     cmd.extend(gpu_args)
 
-    container_cmd = (
-        "cd /jax && %s" % test_cmd
-    )
+    container_cmd = "cd /jax && %s" % test_cmd
 
     cmd.append(image_name)
     cmd.extend(
@@ -229,7 +227,9 @@ def parse_args():
     bdp = subp.add_parser("build_dockers")
 
     testp = subp.add_parser("test")
-    testp.add_argument("--test-cmd", help="Command which will be run inside the test container")
+    testp.add_argument(
+        "--test-cmd", help="Command which will be run inside the test container"
+    )
     testp.add_argument("image_name")
 
     return p.parse_args()

stack.py

@@ -112,7 +112,6 @@ def setup_development(jax_ref: str, xla_ref: str, rebuild_makefile: bool = False
             mf.write(makefile_content)
 
 
-
 def dev_docker():
     cur_abs_path = os.path.abspath(os.curdir)
     image_name = "ubuntu:22.04"
@@ -128,17 +127,20 @@ def dev_docker():
         "--device=/dev/dri",
         "--ipc=host",
         "--shm-size=16G",
-        "--group-add", "video",
+        "--group-add",
+        "video",
         "--cap-add=SYS_PTRACE",
-        "--security-opt", "seccomp=unconfined",
+        "--security-opt",
+        "seccomp=unconfined",
         "-v",
         "%s:/rocm-jax" % cur_abs_path,
-        "--env", "ROCM_JAX_DIR=/rocm-jax",
-        "--env", "_IS_ENTRYPOINT=1",
-        "--entrypoint=%s" % ep
+        "--env",
+        "ROCM_JAX_DIR=/rocm-jax",
+        "--env",
+        "_IS_ENTRYPOINT=1",
+        "--entrypoint=%s" % ep,
     ]
 
-
     cmd.append(image_name)
 
     p = subprocess.Popen(cmd)
@@ -147,6 +149,7 @@ def dev_docker():
 
 # build mode setup
 
+
 # install jax/jaxlib from known versions
 # setup build/install/test script
 def setup_build():
@@ -159,14 +162,27 @@ def parse_args():
     subp = p.add_subparsers(dest="action", required=True)
 
     dev = subp.add_parser("develop")
-    dev.add_argument("--rebuild-makefile", help="Force rebuild of Makefile from template.", action="store_true")
-    dev.add_argument("--xla-ref", help="XLA commit reference to checkout on clone", default=XLA_REPO_REF)
-    dev.add_argument("--jax-ref", help="JAX commit reference to checkout on clone", default=JAX_REPO_REF)
+    dev.add_argument(
+        "--rebuild-makefile",
+        help="Force rebuild of Makefile from template.",
+        action="store_true",
+    )
+    dev.add_argument(
+        "--xla-ref",
+        help="XLA commit reference to checkout on clone",
+        default=XLA_REPO_REF,
+    )
+    dev.add_argument(
+        "--jax-ref",
+        help="JAX commit reference to checkout on clone",
+        default=JAX_REPO_REF,
+    )
 
     docker = subp.add_parser("docker")
 
     return p.parse_args()
 
+
 def main():
     args = parse_args()
     if args.action == "docker":

tools/blacken.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+FILES=(
+  build/ci_build
+  stack.py
+)
+
+FILES+=($(find build -name "*.py"))
+FILES+=($(find tools -name "*.py"))
+
+black -t py36 $* ${FILES[@]}