Repository: pytorch
Issue: Better syntax for OpInfo
Labels: module: tests,triaged,better-engineering

This proposal is loosely related to #55152 but I preferred to split them because I think this one is more controversial.

The `op_db` database is getting bigger and bigger and it's becoming hard to navigate inside `common_method_invocations.py`. In particular, one sub-optimal aspect is that the definition of the `sample_inputs` functions are far away from the actual `OpInfo` call. One possible solution is to completely change the syntax, e.g. something along these lines:

```python
class flip(OpInfo):
    op = torch.flip
    dtypes = all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16)
    supports_out = False

    def sample_inputs(self):
        # this can also be a generator, might be easier to yield multiple 
        # samples instead of manually accumulating them into a list
        return [...]
```

The advantage of keeping `sample_inputs` close to the rest of `OpInfo` is obvious, and it will become even more important if in the future we will ever need to introduce other functions like that. E.g. we already have `ref` which often requires a `lambda`.

Moreover, a class-based opdb is more flexible. Consider the following existing code, which contains a lot of duplication:
```python
    OpInfo('div',
           variant_test_name='true_rounding',
           dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16),
           sample_inputs_func=partial(sample_inputs_div, rounding_mode='true'),
           skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),),
           assert_autodiffed=True),
    OpInfo('div',
           variant_test_name='trunc_rounding',
           dtypes=all_types_and(torch.half, torch.bfloat16),
           sample_inputs_func=partial(sample_inputs_div, rounding_mode='trunc'),
           skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),),
           assert_autodiffed=True),
    OpInfo('div',
           variant_test_name='floor_rounding',
           dtypes=all_types_and(torch.half, torch.bfloat16),
           sample_inputs_func=partial(sample_inputs_div, rounding_mode='floor'),
           skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),),
           assert_autodiffed=True),
```

With class-based syntax it might become something along these lines:
```python
class div_base(OpInfo):
    op_name = 'div'
    rounding_mode = None
    skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),),
    assert_autodiffed=True
    def sample_inputs(self):
        yield SampleInput(..., rounding_mode=self.rounding_mode)

class div_true_rounding(div_base):
    rounding_mode = 'true'
    dtypes = all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16)

class div_trunc_rounding(div_base):
    rounding_mode = 'trunc'
    dtypes=all_types_and(torch.half, torch.bfloat16)

class div_floor_rounding(div_base):
    rounding_mode = 'floow'
    dtypes=all_types_and(torch.half, torch.bfloat16)
```

An open question is how to "register" the OpInfo subclasses inside the db. There are at least two options:
1. use a metaclass, e.g. something along these lines (but engineered better :)):
```python
class OpInfoMeta(type):
    def __new__(self, name, bases, dic):
        ...
        op_db.append(self())

class OpInfo(metaclass=OpInfoMeta):
    ...
```

2. use some kind of class decorator:
```python
@op_db.register
class div(OpInfo):
    ...
```

Solution 2 is probably my preferred because it's less magical and it makes it trivial to define multiple `op_db` databases (and/or attach the same `OpInfo` to multiple databases). It still probably requires a small bit of magic to ensure that we don't forget the `@some_db.register` decorator but this should be easy.

Finally, another advantage of this scheme is that it makes it easier to split the `OpInfo` in multiple files, especially considering that `common_method_invocations.py` is almost 5000 lines by now. For example, I am thinking of something like this:

> #### `common_methods_invocation.py`
``` python
...
op_db = OpDb()
import .opinfo_torch # side effects
import .opinfo_torch_linalg # side effects
```

> #### `op_info_torch.py`
``` python
...
@op_db.register
class add(OpInfo):
    ...

@op_db.register
class div(OpInfo)
    ...
```

> #### `op_info_torch_linalg.py`
``` python
...
@op_db.register
class linalg_det(OpInfo):
    ...

@op_db.register
class linalg_svd(OpInfo)
    ...
```

If there is interest in such a thing, I'd be happy to work on it.

/cc @mruberry @VitalyFedyunin @walterddr