Repository: pytorch
Issue: Better handling of OpInfo.sample_inputs
Labels: high priority,module: tests,triaged,better-engineering,needs design,ezyang's list

## 🚀 Feature
I'd like to propose an `OpInfo` refactoring to enhance the way `sample_inputs` are handled. In particular, I'd like error messages to include the name and/or index of the sample (which was called `variant` in the old `method_tests()` jargon).

## Motivation

Most `OpInfo` have multiple sample inputs and it happens often that a certain opinfo-based test fails only on a subset of samples. Knowing which ones fail is very useful to give a first-guess of what is the problem, but at the moment it is impossible to know by looking at the traceback. This would be especially handy for tests which pass on the development machines but fail on more exotic configurations such as ROCM.

Currently, the only way to know if a test fails only on a subset of samples is to manually comment/uncomment them from the `sample_inputs` function, and to manually keep track of failures. And the same happens if you want to run the test only on a specific sample. This is obviously error-prone and time consuming.

## Pitch

There are at least two options to achieve this:

1. introduce a `sample` fixture, similar to `device` and `dtype`. So, instead of writing:
```python
@ops(op_db)
def test_out(self, device, dtype, op):
    for sample in op.sample_inputs(device, dtype):
        check(sample)
```

you would write
```python
@ops(op_db)
def test_out(self, device, dtype, op, sample):
    check(sample)
```

There are also tests which are doing something only on the *first* sample. For them, we could introduce a fixture named `sample0` or similar.

2. Leave things as they are, but make `sample` a context manager, so that it has the chance to intercept errors and prefix its name to the error:
```python
@ops(op_db)
def test_out(self, device, dtype, op):
    for sample in op.sample_inputs(device, dtype):
        with sample:
            check(sample)
```

Implementing the context manager is as easy as this:
```python
class SampleInput(object):
    ...
    def __enter__(self):
        return self

    def __exit__(self, etype, evalue, tb):
        if etype and self.name:
            evalue.args = (f'[SampleInput: {self.name}] {evalue}',)
```

I think that solution 1 is better in all aspects, but I wanted to mentioned 2 for completeness and/or in case I missed anything. The only two "non-problems" of solution 1 which I can imagine are:

* proliferation of tests: we will generate many more tests than now, but I don't see how/why it would be a problem. Also, I think it's what happened before with `method_tests()`

* many individual tests could be slower than a single test, in case there is some expensive setup/teardown logic before/after the `for sample in op.sample_inputs()`. At a quick glance I don't think we have any test like that at the moment, but in any case nothing would prevent us not to use the `sample` fixture and do the `for` loop manually.

If there is interest in such a feature, I'd be happy to work on that.

/cc @ezyang @gchanan @zou3519 @mruberry @VitalyFedyunin @walterddr