[test] [refactor] Fix test_element_wise.py long time (#1120)

* [test] [refactor] Fix test_element_wise.py long time

* [skip ci] fix ident

* [skip ci] Update python/taichi/main.py

Co-authored-by: Chengchen(Rex) Wang <14366016+rexwangcc@users.noreply.github.com>

* [skip ci] fix missing prog=

* [skip ci] enforce code format

* fix test_cli

* [skip ci] enforce code format

* fix test_minmax

* [skip ci] enforce code format

Co-authored-by: Chengchen(Rex) Wang <14366016+rexwangcc@users.noreply.github.com>
Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

python/taichi/__init__.py

@@ -8,7 +8,7 @@
 from taichi.misc.gui import GUI
 from taichi.misc.image import imread, imwrite, imshow
 from taichi.misc.task import Task
-from taichi.misc.test import approx
+from taichi.misc.test import *
 from taichi.misc import settings as settings
 from taichi.misc.gui import rgb_to_hex
 from taichi.misc.settings import *

python/taichi/lang/__init__.py

@@ -452,6 +452,7 @@ def torch_test(func):
 def host_arch_only(func):
     import taichi as ti
 
+    @functools.wraps(func)
     def test(*args, **kwargs):
         archs = [ti.core.host_arch()]
         for arch in archs:

python/taichi/lang/common_ops.py

@@ -151,7 +151,7 @@ def __ior__(self, other):
         self.atomic_or(other)
         return self
 
-    # we don't support atomic_mul/truediv/floordiv yet:
+    # we don't support atomic_mul/truediv/floordiv/mod yet:
     def __imul__(self, other):
         import taichi as ti
         self.assign(ti.mul(self, other))
@@ -167,6 +167,11 @@ def __ifloordiv__(self, other):
         self.assign(ti.floordiv(self, other))
         return self
 
+    def __imod__(self, other):
+        import taichi as ti
+        self.assign(ti.mod(self, other))
+        return self
+
     def assign(self, other):
         import taichi as ti
         return ti.assign(self, other)
@@ -182,6 +187,8 @@ def augassign(self, x, op):
             self /= x
         elif op == 'FloorDiv':
             self //= x
+        elif op == 'Mod':
+            self %= x
         elif op == 'BitAnd':
             self &= x
         elif op == 'BitOr':

python/taichi/lang/matrix.py

@@ -150,8 +150,8 @@ def element_wise_binary(self, foo, other):
             other = Matrix(other)
         if foo.__name__ == 'assign' and not isinstance(other, Matrix):
             raise SyntaxError(
-                f'cannot assign scalar expr to '
-                'taichi class {type(a)}, maybe you want to use `a.fill(b)` instead?'
+                'cannot assign scalar expr to '
+                f'taichi class {type(a)}, maybe you want to use `a.fill(b)` instead?'
             )
         if isinstance(other, Matrix):
             assert self.m == other.m and self.n == other.n

python/taichi/lang/ops.py

@@ -331,6 +331,11 @@ def bit_and(a, b):
     return Expr(taichi_lang_core.expr_bit_and(a.ptr, b.ptr), tb=stack_info())
 
 
+@binary
+def bit_xor(a, b):
+    return Expr(taichi_lang_core.expr_bit_xor(a.ptr, b.ptr), tb=stack_info())
+
+
 # We don't have logic_and/or instructions yet:
 logical_or = bit_or
 logical_and = bit_and
@@ -393,16 +398,9 @@ def ti_max(*args):
     if num_args == 1:
         return args[0]
     elif num_args == 2:
-        if isinstance(args[0], numbers.Number) and isinstance(
-                args[1], numbers.Number):
-            return max(args[0], args[1])
-        else:
-            return Expr(
-                taichi_lang_core.expr_max(
-                    Expr(args[0]).ptr,
-                    Expr(args[1]).ptr))
+        return max(args[0], args[1])
     else:
-        return ti_max(args[0], ti_max(*args[1:]))
+        return max(args[0], ti_max(*args[1:]))
 
 
 def ti_min(*args):
@@ -411,16 +409,9 @@ def ti_min(*args):
     if num_args == 1:
         return args[0]
     elif num_args == 2:
-        if isinstance(args[0], numbers.Number) and isinstance(
-                args[1], numbers.Number):
-            return min(args[0], args[1])
-        else:
-            return Expr(
-                taichi_lang_core.expr_min(
-                    Expr(args[0]).ptr,
-                    Expr(args[1]).ptr))
+        return min(args[0], args[1])
     else:
-        return ti_min(args[0], ti_min(*args[1:]))
+        return min(args[0], ti_min(*args[1:]))
 
 
 def ti_any(a):

python/taichi/main.py

@@ -144,7 +144,8 @@ def example(self, arguments: list = sys.argv[2:]):
         """Run an example by name (or name.py)"""
         choices = TaichiMain._get_available_examples()
 
-        parser = argparse.ArgumentParser(description=f"{self.example.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti example',
+                                         description=f"{self.example.__doc__}")
         parser.add_argument(
             "name",
             help=f"Name of an example (supports .py extension too)\n",
@@ -166,7 +167,8 @@ def example(self, arguments: list = sys.argv[2:]):
     @register
     def release(self, arguments: list = sys.argv[2:]):
         """Make source code release"""
-        parser = argparse.ArgumentParser(description=f"{self.release.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti release',
+                                         description=f"{self.release.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -203,7 +205,8 @@ def _mp4_file(name: str) -> str:
     @register
     def gif(self, arguments: list = sys.argv[2:]):
         """Convert mp4 file to gif in the same directory"""
-        parser = argparse.ArgumentParser(description=f"{self.gif.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti gif',
+                                         description=f"{self.gif.__doc__}")
         parser.add_argument('-i',
                             '--input',
                             required=True,
@@ -230,7 +233,7 @@ def gif(self, arguments: list = sys.argv[2:]):
     def video_speed(self, arguments: list = sys.argv[2:]):
         """Speed up video in the same directory"""
         parser = argparse.ArgumentParser(
-            description=f"{self.video_speed.__doc__}")
+            prog='ti video_speed', description=f"{self.video_speed.__doc__}")
         parser.add_argument('-i',
                             '--input',
                             required=True,
@@ -260,7 +263,7 @@ def video_speed(self, arguments: list = sys.argv[2:]):
     def video_crop(self, arguments: list = sys.argv[2:]):
         """Crop video in the same directory"""
         parser = argparse.ArgumentParser(
-            description=f"{self.video_crop.__doc__}")
+            prog='ti video_crop', description=f"{self.video_crop.__doc__}")
         parser.add_argument('-i',
                             '--input',
                             required=True,
@@ -303,7 +306,7 @@ def video_crop(self, arguments: list = sys.argv[2:]):
     def video_scale(self, arguments: list = sys.argv[2:]):
         """Scale video resolution in the same directory"""
         parser = argparse.ArgumentParser(
-            description=f"{self.video_scale.__doc__}")
+            prog='ti video_scale', description=f"{self.video_scale.__doc__}")
         parser.add_argument('-i',
                             '--input',
                             required=True,
@@ -343,7 +346,8 @@ def video_scale(self, arguments: list = sys.argv[2:]):
     @register
     def video(self, arguments: list = sys.argv[2:]):
         """Make a video using *.png files in the current directory"""
-        parser = argparse.ArgumentParser(description=f"{self.video.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti video',
+                                         description=f"{self.video.__doc__}")
         parser.add_argument("inputs", nargs='*', help="PNG file(s) as inputs")
         parser.add_argument('-o',
                             '--output',
@@ -377,7 +381,8 @@ def video(self, arguments: list = sys.argv[2:]):
     @register
     def doc(self, arguments: list = sys.argv[2:]):
         """Build documentation"""
-        parser = argparse.ArgumentParser(description=f"{self.doc.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti doc',
+                                         description=f"{self.doc.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -390,7 +395,8 @@ def doc(self, arguments: list = sys.argv[2:]):
     def update(self, arguments: list = sys.argv[2:]):
         """Update the Taichi codebase"""
         # TODO: Test if this still works, fix if it doesn't
-        parser = argparse.ArgumentParser(description=f"{self.update.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti update',
+                                         description=f"{self.update.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -401,7 +407,8 @@ def update(self, arguments: list = sys.argv[2:]):
     @register
     def format(self, arguments: list = sys.argv[2:]):
         """Reformat modified source files"""
-        parser = argparse.ArgumentParser(description=f"{self.format.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti format',
+                                         description=f"{self.format.__doc__}")
         parser.add_argument(
             'diff',
             nargs='?',
@@ -417,7 +424,7 @@ def format(self, arguments: list = sys.argv[2:]):
     def format_all(self, arguments: list = sys.argv[2:]):
         """Reformat all source files"""
         parser = argparse.ArgumentParser(
-            description=f"{self.format_all.__doc__}")
+            prog='ti format_all', description=f"{self.format_all.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -427,7 +434,8 @@ def format_all(self, arguments: list = sys.argv[2:]):
     @register
     def build(self, arguments: list = sys.argv[2:]):
         """Build C++ files"""
-        parser = argparse.ArgumentParser(description=f"{self.build.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti build',
+                                         description=f"{self.build.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -545,7 +553,7 @@ def _get_benchmark_output_dir():
     def regression(self, arguments: list = sys.argv[2:]):
         """Display benchmark regression test result"""
         parser = argparse.ArgumentParser(
-            description=f"{self.regression.__doc__}")
+            prog='ti regression', description=f"{self.regression.__doc__}")
         parser.add_argument('files',
                             nargs='*',
                             help='Test file(s) to be run for benchmarking')
@@ -568,7 +576,7 @@ def regression(self, arguments: list = sys.argv[2:]):
     def baseline(self, arguments: list = sys.argv[2:]):
         """Archive current benchmark result as baseline"""
         parser = argparse.ArgumentParser(
-            description=f"{self.baseline.__doc__}")
+            prog='ti baseline', description=f"{self.baseline.__doc__}")
         args = parser.parse_args(arguments)
 
         # Short circuit for testing
@@ -654,7 +662,7 @@ def _test_cpp(args):
     def benchmark(self, arguments: list = sys.argv[2:]):
         """Run Python tests in benchmark mode"""
         parser = argparse.ArgumentParser(
-            description=f"{self.benchmark.__doc__}")
+            prog='ti benchmark', description=f"{self.benchmark.__doc__}")
         parser.add_argument('files', nargs='*', help='Test file(s) to be run')
         parser.add_argument('-T',
                             '--tprt',
@@ -708,7 +716,8 @@ def benchmark(self, arguments: list = sys.argv[2:]):
     @register
     def test(self, arguments: list = sys.argv[2:]):
         """Run the tests"""
-        parser = argparse.ArgumentParser(description=f"{self.test.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti test',
+                                         description=f"{self.test.__doc__}")
         parser.add_argument('files',
                             nargs='*',
                             help='Test name(s) to be run, e.g. "cli"')
@@ -737,8 +746,34 @@ def test(self, arguments: list = sys.argv[2:]):
             dest='threads',
             type=str,
             help='Custom number of threads for parallel testing')
+        parser.add_argument(
+            '-a',
+            '--arch',
+            required=False,
+            default=None,
+            dest='arch',
+            type=str,
+            help='Custom the arch(s) (backend) to run tests on')
+        parser.add_argument(
+            '-n',
+            '--exclusive',
+            required=False,
+            default=False,
+            dest='exclusive',
+            action='store_true',
+            help=
+            'Exclude arch(s) from test instead of include them, together with -a'
+        )
+
         args = parser.parse_args(arguments)
 
+        if args.arch:
+            arch = args.arch
+            if args.exclusive:
+                arch = '^' + arch
+            print(f'Running on Arch={arch}')
+            os.environ['TI_WANTED_ARCHS'] = arch
+
         # Short circuit for testing
         if self.test_mode: return args
 
@@ -760,7 +795,8 @@ def test(self, arguments: list = sys.argv[2:]):
     @register
     def debug(self, arguments: list = sys.argv[2:]):
         """Debug a single script"""
-        parser = argparse.ArgumentParser(description=f"{self.debug.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti debug',
+                                         description=f"{self.debug.__doc__}")
         parser.add_argument(
             'filename',
             help='A single (Python) script to run with debugger, e.g. render.py'
@@ -781,7 +817,8 @@ def debug(self, arguments: list = sys.argv[2:]):
     @register
     def run(self, arguments: list = sys.argv[2:]):
         """Run a specific task"""
-        parser = argparse.ArgumentParser(description=f"{self.run.__doc__}")
+        parser = argparse.ArgumentParser(prog='ti run',
+                                         description=f"{self.run.__doc__}")
         parser.add_argument('taskname',
                             help='A single task name to run, e.g. test_math')
         parser.add_argument('taskargs',

python/taichi/misc/test.py

@@ -1,21 +1,25 @@
-import taichi as ti
-
+def approx(expected, **kwargs):
+    import taichi as ti
+    from pytest import approx
 
-class boolean_integer:
-    def __init__(self, value):
-        self.value = value
+    class boolean_integer:
+        def __init__(self, value):
+            self.value = value
 
-    def __eq__(self, other):
-        return bool(self.value) == bool(other)
+        def __eq__(self, other):
+            return bool(self.value) == bool(other)
 
+        def __ne__(self, other):
+            return bool(self.value) != bool(other)
 
-def approx(expected, **kwargs):
     if isinstance(expected, bool):
         return boolean_integer(expected)
 
-    from pytest import approx
-
     if ti.cfg.arch == ti.opengl:
         kwargs['rel'] = max(kwargs.get('rel', 1e-6), 1e-3)
 
     return approx(expected, **kwargs)
+
+
+def allclose(x, y, **kwargs):
+    return x == approx(y, **kwargs)

taichi/python/export_lang.cpp

@@ -374,6 +374,7 @@ void export_lang(py::module &m) {
   m.def("expr_add", expr_add);
   m.def("expr_sub", expr_sub);
   m.def("expr_mul", expr_mul);
+  m.def("expr_div", expr_div);
   m.def("expr_truediv", expr_truediv);
   m.def("expr_floordiv", expr_floordiv);
   m.def("expr_mod", expr_mod);

tests/python/test_cli.py

@@ -172,8 +172,7 @@ def test_cli_video():
 
 
 def test_cli_format():
-    with patch_sys_argv_helper(["ti", "format", "-d",
-                                "51fff7af"]) as custom_argv:
+    with patch_sys_argv_helper(["ti", "format", "51fff7af"]) as custom_argv:
         cli = TaichiMain(test_mode=True)
         args = cli()
         assert args.diff == "51fff7af"