Add regression test for torch.expand (#3785)

Author: ymwangg

.circleci/common.sh

@@ -58,11 +58,11 @@ function install_deps_pytorch_xla() {
   pip install hypothesis
   pip install cloud-tpu-client
   pip install absl-py
-  pip install --upgrade numpy>=1.18.5
+  pip install --upgrade "numpy>=1.18.5"
   pip install --upgrade numba
 
   # Using the Ninja generator requires CMake version 3.13 or greater
-  pip install cmake>=3.13 --upgrade
+  pip install "cmake>=3.13" --upgrade
 
   sudo apt-get -qq update
 

.circleci/docker/install_conda.sh

@@ -39,7 +39,7 @@ function install_and_setup_conda() {
   /usr/bin/yes | pip install cloud-tpu-client
   /usr/bin/yes | pip install expecttest==0.1.3
   /usr/bin/yes | pip install ninja  # Install ninja to speedup the build
-  /usr/bin/yes | pip install cmake>=3.13 --upgrade  # Using Ninja requires CMake>=3.13
+  /usr/bin/yes | pip install "cmake>=3.13" --upgrade  # Using Ninja requires CMake>=3.13
   /usr/bin/yes | pip install absl-py
   # Additional PyTorch requirements
   /usr/bin/yes | pip install scikit-image scipy==1.1.0  # >1.1.0 breaks PyTorch tests

test/run_tests.sh

@@ -111,6 +111,7 @@ function run_op_tests {
   run_xla_ir_debug python3 "$CDIR/test_env_var_mapper.py"
   run_pjrt python3 "$CDIR/pjrt/test_experimental_pjrt.py"
   run_pjrt python3 "$CDIR/pjrt/test_experimental_tpu.py"
+  run_test python3 "$CDIR/test_operations_hlo.py" "$@" --verbosity=$VERBOSITY
 }
 
 function run_mp_op_tests {

test/test_operations_hlo.py

@@ -0,0 +1,47 @@
+# Parse local options first, and rewrite the sys.argv[].
+# We need to do that before import "common", as otherwise we get an error for
+# unrecognized arguments.
+import argparse
+import sys
+
+parser = argparse.ArgumentParser(add_help=False)
+parser.add_argument('--replicated', action='store_true')
+parser.add_argument('--long_test', action='store_true')
+parser.add_argument('--max_diff_count', type=int, default=25)
+parser.add_argument('--verbosity', type=int, default=0)
+FLAGS, leftovers = parser.parse_known_args()
+sys.argv = [sys.argv[0]] + leftovers
+
+# Normal imports section starts here.
+import torch
+import torch_xla
+import torch_xla.utils.utils as xu
+import torch_xla.core.xla_model as xm
+import torch_xla.debug.metrics as met
+import unittest
+
+
+class TestOperationsHlo(unittest.TestCase):
+
+  def setUp(self):
+    super(TestOperationsHlo, self).setUp()
+
+  def tearDown(self):
+    super(TestOperationsHlo, self).tearDown()
+
+  def test_expand(self):
+    a = torch.rand(1, 5, device=xm.xla_device())
+    b = a.expand(5, 5)
+    hlo_text = torch_xla._XLAC._get_xla_tensors_text([b])
+    assert 'aten::expand' in hlo_text
+
+
+if __name__ == '__main__':
+  torch.set_default_tensor_type('torch.FloatTensor')
+  torch.manual_seed(42)
+  torch_xla._XLAC._xla_set_use_full_mat_mul_precision(
+      use_full_mat_mul_precision=True)
+  test = unittest.main(verbosity=FLAGS.verbosity, exit=False)
+  if xu.getenv_as('METRICS_DEBUG', bool, defval=False):
+    print(met.metrics_report())
+  sys.exit(0 if test.result.wasSuccessful() else 1)