Merge branch 'main' into gh/swolchok/517/head

Author: swolchok

.github/workflows/pull.yml

@@ -632,32 +632,33 @@ jobs:
         # run eval_llama wikitext task
         PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_wikitext.sh
 
-  test-eval_llama-mmlu-linux:
-    name: test-eval_llama-mmlu-linux
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
-    permissions:
-      id-token: write
-      contents: read
-    strategy:
-      fail-fast: false
-    with:
-      runner: linux.24xlarge
-      docker-image: ci-image:executorch-ubuntu-22.04-clang12
-      submodules: 'recursive'
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
-      script: |
-        # The generic Linux job chooses to use base env, not the one setup by the image
-        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
-        conda activate "${CONDA_ENV}"
-
-        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool "cmake"
-
-        # install llama requirements
-        bash examples/models/llama/install_requirements.sh
-
-        # run eval_llama mmlu task
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh
+  # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2
+  # test-eval_llama-mmlu-linux:
+  #   name: test-eval_llama-mmlu-linux
+  #   uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+  #   permissions:
+  #     id-token: write
+  #     contents: read
+  #   strategy:
+  #     fail-fast: false
+  #   with:
+  #     runner: linux.24xlarge
+  #     docker-image: ci-image:executorch-ubuntu-22.04-clang12
+  #     submodules: 'recursive'
+  #     ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+  #     timeout: 90
+  #     script: |
+  #       # The generic Linux job chooses to use base env, not the one setup by the image
+  #       CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+  #       conda activate "${CONDA_ENV}"
+
+  #       PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool "cmake"
+
+  #       # install llama requirements
+  #       bash examples/models/llama/install_requirements.sh
+
+  #       # run eval_llama mmlu task
+  #       PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh
 
   test-llama_runner_eager-linux:
     name: test-llama_runner_eager-linux

exir/backend/test/test_backends.py

@@ -1033,7 +1033,7 @@ def false_fn(x, y):
 
         def f(x, y):
             x = x + y
-            x = torch.ops.higher_order.cond(x[0][0] == 1, true_fn, false_fn, [x, y])
+            x = torch.cond(x[0][0] == 1, true_fn, false_fn, [x, y])
             x = x - y
             return x
 

exir/tests/control_flow_models.py

@@ -20,9 +20,7 @@ def true_branch(x):
         def false_branch(x):
             return x * x
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_random_inputs(self):
         return (torch.rand(5),)
@@ -39,9 +37,7 @@ def true_branch(x):
         def false_branch(x):
             return x * x * x
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_upper_bound_inputs(self):
         return (torch.rand(8),)
@@ -72,9 +68,7 @@ def true_branch(x):
         def false_branch(x):
             return x * 2
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_random_inputs(self):
         return (torch.eye(5) * 2,)

exir/tests/test_passes.py

@@ -1463,9 +1463,7 @@ def forward(self, pred, x):
                 out = torch.nn.functional.linear(
                     x, self.w.to(torch.float16).to(torch.float32)
                 )
-                return torch.ops.higher_order.cond(
-                    pred, self.true_fn, self.false_fn, [out]
-                )
+                return torch.cond(pred, self.true_fn, self.false_fn, [out])
 
         mod = Module()
         x = torch.randn([3, 3])