Make test_types.py work for the Metal backend

[k-ye]

Related issue id = #504 #396

Hi @archibate

Sorry i changed your tests quite a bit, in order to make it work for the Metal backend as well. The problem with Metal is that it doesn't support 64bit buffers, so we had to separate the tests into 64-bit vs the rest.

I found a nice decorator @pytest.mark.parametrize() (doc) that can somewhat simplify the code.

Also I saw that you got this "variables are expected to be declared before kernel invocation" problem, therefore had to create another wrapper inside each test? IIUC, ti.all_archs(_with) resets the test case here:

taichi/python/taichi/lang/__init__.py

Line 224 in 57287cf

ti.init(arch=arch, **kwargs)

So by putting @ti.all_archs at the bottom of the decorator chain, we no longer need to wrap it, because it resets Taichi every time the test case is called. (Unfortunately, this doesn't apply to this PR, as explained below)

At last, I tried something like this:

def _test_foo(x):
  ...

@pytest.mark.parametrize('x', [1, 2, 3])
@ti.all_archs
def test_foo(x):
  _test_foo(x)

But @pytest.mark.parametrize didn't work well with @ti.all_archs this way. So I had to once again, wrap the _test_foo() inside test_foo(), so that @ti.all_archs can correctly apply the decoration:

@pytest.mark.parametrize('x', [1, 2, 3])
def test_foo(x):
  @ti.all_archs
  def run():
    _test_foo(x)
  run()

[archibate]

Also I saw that you got this "variables are expected to be declared before kernel invocation" problem, therefore had to create another wrapper inside each test? IIUC, ti.all_archs(_with) resets the test case.

Sounds to be a good solution for this!

But @pytest.mark.parametrize didn't work well with @ti.all_archs this way.

So that the problem is, ti.all_archs didn't pass x to test_foo in the first case?

[archibate]

That's strange, in:

taichi/python/taichi/lang/__init__.py

Line 211 in 57287cf

def wrapped(*test_args, **test_kwargs):

test arguments should be recognized...

[k-ye]

That's strange, in:

taichi/python/taichi/lang/__init__.py

Line 211 in 57287cf

def wrapped(*test_args, **test_kwargs):

test arguments should be recognized...

Yeah, I think pytest.mark.parametrize only knows about the normal function arguments, no defaults, no *args, no **kwargs... pytest-dev/pytest#3221 Unfortunately we have to put @ti.all_archs at the bottom of this chain :(

Similar issue: https://stackoverflow.com/questions/12197263/pytest-parameterize-decorator-not-working-with-kwargs

[archibate]

Maybe we could customize our own parametrize like:

def parametrize(xs):
  def decorator(foo):
    def wrapped(*args, **kwargs):
      for x in xs:
        foo(x, *args, **kwargs)
    return wrapped
  return decorator

@parametrize([1, 2, 3])
@ti.all_archs
def test_func(x):
  assert x == 233

[k-ye]

Maybe we can customize our own parametrize like

SG! While I generally don't prefer maintaining our own utilities for such things that are not essential to Taichi, I guess having a @ti.parametrize will work better for our tests, because most of them need to use @ti.all_arch. Thanks!

[archibate]

Cool implementation! I just opened an issue for this, see pytest-dev/pytest#6810.

[archibate]

Their reply: pytest-dev/pytest#6810 (comment).
They say we can solve this by adding @functools.wraps before this line:

taichi/python/taichi/lang/__init__.py

Line 211 in 57287cf

def wrapped(*test_args, **test_kwargs):

How do you think?

[k-ye]

Thank you so much for following up with the pytest team! Yeah, after decorating the tests with @functools.wraps, we can now integrate with @pytest.mark.parametrize!! Again, really appreciate your efforts! (BTW I'm not available during working hours, so I will have to catch up new comments later tonight..)

[yuanming-hu]

Looks good to me. Thank you @k-ye @archibate!