Repository: pytorch
Issue: test_variant_consistency_jit fails for torch.tensordot with dtype float32 with error INTERNAL ASSERT FAILED
Labels: oncall: jit,module: tests

## 🐛 Bug

`test_variant_consistency_jit` fails for `tensordot` with dtype `torch.float32` on both CPU & CUDA.
This test passes for `torch.complex64`, though.

`tensordot` has been _jittable_ since #53672.
https://github.com/pytorch/pytorch/blob/master/docs/source/jit_unsupported.rst states that JIT is unsupported for `tensordot`, but it should be modified (I had closed this issue on its basis, initially).

This test fails for `float32` because of the following call to `check_alias_annotation`, whose invocation doesn't occur for `complex64`:
https://github.com/pytorch/pytorch/blob/5cab3b9cf628ba44b599fa7da4c06404986dd656/test/test_ops.py#L397-L403

So, maybe this test fails for `tensordot` with `float32` because `tensordot`'s a functional.

## To Reproduce

Steps to reproduce the behavior:

1. Create an OpInfo for `tensordot` with some sample inputs. An example is in the Additional context section:
.
2. Run `python test/test_ops.py TestCommonCPU.test_variant_consistency_jit_tensordot_cpu_float32`
or `python test/test_ops.py TestCommonCUDA.test_variant_consistency_jit_tensordot_cuda_float32`.

3. The test fails with the following error:
```
ERROR: test_variant_consistency_jit_tensordot_cuda_float32 (__main__.TestCommonCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/torch/testing/_internal/common_utils.py", line 1030, in wrapper
    method(*args, **kwargs)
  File "/torch/testing/_internal/common_utils.py", line 1030, in wrapper
    method(*args, **kwargs)
  File "/torch/testing/_internal/common_device_type.py", line 297, in instantiated_test
    raise rte
  File "/torch/testing/_internal/common_device_type.py", line 292, in instantiated_test
    result = test_fn(self, *args)
  File "/torch/testing/_internal/common_device_type.py", line 266, in test_wrapper
    return test(*args, **kwargs)
  File "test_ops.py", line 403, in test_variant_consistency_jit
    func_type=func_type, aten_name=op.aten_name)
  File "/torch/testing/_internal/jit_metaprogramming_utils.py", line 486, in check_alias_annotation
    torch._C._jit_check_alias_annotation(CU.the_method.graph, tuple(tensors), aten_name)
RuntimeError: aliasOp != torch::jit::getOperatorAliasMap().end()INTERNAL ASSERT FAILED at 
"/torch/csrc/jit/passes/utils/check_alias_annotation.cpp":155, please report a bug to PyTorch.
```

## Expected behavior

The test should pass on both CPU & CUDA, as it does for `complex64`.

Maybe `test_variant_consistency_jit` shouldn't call `check_alias_annotation` for functionals such as `tensordot`.

## Environment
OpInfo added on top of current master branch's source-code, as in #56286.
This issue occurs even if the OpInfo is added without any changes to the `tensordot` source code,


## Additional context

OpInfo & its sample inputs can be added in `torch/testing/_internal/common_methods_invocations.py`.

```
    OpInfo('tensordot',
           dtypes=floating_and_complex_types_and(torch.half),
           dtypesIfCPU=all_types_and_complex_and(torch.half, torch.bfloat16),
           dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []),
           dtypesIfROCM=floating_and_complex_types_and(torch.half, torch.bfloat16),
           supports_autograd=False,
           sample_inputs_func=sample_inputs_tensordot)
```
```
def sample_inputs_tensordot(self, device, dtype, requires_grad, **kwargs):
    return (
        SampleInput(
            make_tensor((S, S, S), device, dtype, requires_grad=requires_grad),
            args=(
                make_tensor((S, S, S), device, dtype, requires_grad=requires_grad),
            ),
            kwargs=dict(dims=2,),
        ),
        SampleInput(
            make_tensor((S, S, 1), device, dtype, requires_grad=requires_grad),
            args=(
                make_tensor((S, 1, S), device, dtype, requires_grad=requires_grad),
            ),
            kwargs=dict(dims=([0, 2], [2, 1]),),
        ),
    )
```

cc @mruberry @VitalyFedyunin @walterddr @gmagogsfm