Update pytorch branch for Terraform setup. (#4889)

* Change PyTorch default to main

* Fix file paths after moving config to a new dir

* Use 'HEAD' for PyTorch and XLA versions in nightly builds

* Checkout at origin/<branch> for the git_checkout step

Author: mateuszlewko

infra/terraform_modules/xla_docker_build/variables.tf

@@ -65,7 +65,7 @@ variable "dockerfile" {
 }
 
 variable "docker_context_dir" {
-  default = "docker/experimental/ansible"
+  default = "infra/ansible"
 }
 
 variable "docker_repo_url" {

infra/terraform_modules/xla_docker_build/xla_docker_build.tf

@@ -24,13 +24,13 @@ locals {
     {
       id   = "git_fetch"
       name = "gcr.io/cloud-builders/git"
-      args = ["fetch", "origin", var.ansible_branch]
+      args = ["fetch", "--unshallow"]
     },
     {
       id   = "git_checkout"
       name = "gcr.io/cloud-builders/git"
-      args = ["checkout", var.ansible_branch]
-    }
+      args = ["checkout", "origin/${var.ansible_branch}"]
+    },
   ]
 
   build_and_push_docker_image_steps = [

infra/tpu-pytorch-releases/cloud_builds.tf

@@ -56,8 +56,8 @@ module "nightly_builds" {
   ansible_vars = merge(each.value, {
     package_version = var.nightly_package_version
     nightly_release = true
-    pytorch_git_rev = "master"
-    xla_git_rev     = "master"
+    pytorch_git_rev = "HEAD"
+    xla_git_rev     = "HEAD"
   })
 
   ansible_branch      = "master"
@@ -78,7 +78,7 @@ module "nightly_builds" {
       : format("CUDA %s", each.value.cuda_version)
     } docker image and corresponding wheels for PyTorch/XLA.",
     "Trigger managed by Terraform setup in",
-    "docker/experimental/tpu-pytorch-releases/cloud_builds.tf."
+    "infra/tpu-pytorch-releases/cloud_builds.tf."
   ])
 
   wheels_dest = "${module.releases_storage_bucket.url}/wheels/${

infra/tpu-pytorch/cloud_builds.tf

@@ -19,7 +19,7 @@ module "dev_image" {
   description = join(" ", [
     "Build development image with TPU support.",
     "Trigger managed by Terraform setup in",
-    "docker/experimental/tpu-pytorch/cloud_builds.tf.",
+    "infra/tpu-pytorch/cloud_builds.tf.",
   ])
 
   build_args = {