The class name should always contain the name of the op that is being
tested. This makes it easy to search for tests for a particular op. Often times
an op will require multiple tests to make sure different paths in the
compilation work as expected. In such cases, it is customary to add extra
information to the class name about what is being tested. In this example, the
op is being tested with a rank-3 tensor as an input.
### `__init__` Method
```python
def __init__(self):
super().__init__()
```
In most tests, the `__init__` method simply calls the `__init__` method of the
`torch.nn.Module` class. However, sometimes this method can be used to
initialize parameters needed in the `forward` method. An example of such a case
is in the [E2E test for Resnet18](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir_e2e_test/test_suite/vision_models.py#L17-L22).
### `@export` and `@annotate_args` Decorators
```python
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
([-1, -1], torch.int64, True),
])
```
The [`@export` decorator](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir_e2e_test/torchscript/annotations.py#L30)
lets the importer know which methods in the class will be public after the
`torch.nn.Module` gets imported into the `torch` dialect. All E2E tests should
have this decorator on the `forward` method.
The [`@annotate_args` decorator](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir_e2e_test/torchscript/annotations.py#L53)
is used to give the importer information about the arguments of the method being
decorated, which can then be propagated further into the IR of the body of the
method. The list of annotations **must** have one annotation for each argument
including the `self` argument. The `self` argument always gets the annotation of
`None`, while the other inputs get an annotation with three fields in the
following order:
1. Shape of input tensor. Use `-1` for dynamic dimensions
2. Dtype of the input tensor
3. Boolean representing whether the input tensor [has value semantics](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/class_annotator.h#L54-L67). This
will always be true for E2E tests, since the [Torch-MLIR backend contract](architecture.md#the-backend-contract) requires all tensors in the
IR to eventually have value semantics.
From the structure of the annotations for the arguments other than the `self`
argument it is clear that only tensor arguments are supported. This means that
if an op requires an input other than a tensor, you need to do one of the
following:
- Create the value in the method body
- Create the value as a class parameter in the `__init__` method
- In the case of certain values such as `int`s and `float`s, you can pass a
zero-rank tensor as an input and use `int(input)` or `float(input)`in the
method body to turn the tensor into a scalar `int` or `float`, respectively.
### `forward` Method
```python
def forward(self, x, index):
return torch.ops.aten.index(x, (index,))
```
The forward method should be a simple test of your op. In other words, it will
almost always take the form of simply returning the result of calling your
op. The call to your op should **always** be made using
`torch.ops.aten.{op_name}` to make it very clear which ATen op is being
tested. Some ATen ops have different variants under the same base name, such as
`aten.mean`, which has also a variant `aten.mean.dim`. At the Python level, such
ops are accessed by just their base name, and the right variant is chosen based
on the inputs given. For example, to test `aten.mean.dim` the test should use
