Update gcp scripts with minor corrections

scripts/gcp/README.md

@@ -48,9 +48,9 @@ $ gcloud beta services identity create --service tpu.googleapis.com
 - Clone HEIR repo and install dependencies
 
 ```sh
-$ git clone git@github.com:google/heir.git
+$ git clone https://github.com/google/heir.git
 $ cd heir
-$ pip install -r /scripts/gcp/requirements.txt
+$ pip install -r ./scripts/gcp/requirements.txt
 ```
 
 - Create a TPU

scripts/gcp/core.py

@@ -50,7 +50,8 @@ def add_args(parser) -> None:
     parser.add_argument("--project", type=str, required=False)
     parser.add_argument("--zone", type=str, required=False)
 
-  def from_args(self, args) -> "Core":
+  @staticmethod
+  def from_args(args) -> "Core":
     """Creates a Core object from the given args."""
     return Core(project=args.project, zone=args.zone)
 

scripts/gcp/provision.py

@@ -128,7 +128,9 @@ def provision(
 #!/bin/bash
 set -eux
 pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install --upgrade jax jaxlib
 pip install jaxite
+pip install absl-py
 """,
           },
           # cidr_block !!!

scripts/gcp/tool

@@ -207,7 +207,7 @@ class Tool:
         "--accelerator-type",
         type=str,
         # https://cloud.google.com/tpu/docs/system-architecture-tpu-vm
-        default="v5litepod-8",
+        default="v5litepod-4",
         help="Accelerator type",
     )
     parser.add_argument(