Repository: pytorch
Issue: TorchScript does not recognize mix-in types with `Enum`
Labels: oncall: jit

### 🐛 Describe the bug

When inheriting str and Enum to create `StrEnum` class (similar to `IntEnum`) https://docs.python.org/3/library/enum.html#others, TorchScript does not recognize that these enums are comparable with `str` types. Below is an example that use `(str, Enum)`, which works fine but cannot be scripted.

```python
class Color(str, Enum):
    RED = "1"
    GREEN = "2"


@torch.jit.script
def test_str_enum(x: str) -> str:
    if x == Color.GREEN.value:
        return "green"
    if x == Color.RED:
        return "red"
```

```
RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor):
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'str'.
  
  aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor):
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'str'.
  
  aten::eq.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> (Tensor(a!)):
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'str'.
  
  aten::eq.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> (Tensor(a!)):
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'str'.
  
  aten::eq.int_list(int[] a, int[] b) -> (bool):
  Expected a value of type 'List[int]' for argument 'a' but instead found type 'str'.
  
  aten::eq.device(Device a, Device b) -> (bool):
  Expected a value of type 'Device' for argument 'b' but instead found type 'Enum<__torch__.Color>'.
  
  aten::eq.bool(bool a, bool b) -> (bool):
  Expected a value of type 'bool' for argument 'a' but instead found type 'str'.
  
  aten::eq.enum(AnyEnumType a, AnyEnumType b) -> (bool):
  Expected a value of type 'AnyEnumType' for argument 'a' but instead found type 'str'.
  
  aten::eq.int(int a, int b) -> (bool):
  Expected a value of type 'int' for argument 'a' but instead found type 'str'.
  
  aten::eq.complex(complex a, complex b) -> (bool):
  Expected a value of type 'complex' for argument 'a' but instead found type 'str'.
  
  aten::eq.float(float a, float b) -> (bool):
  Expected a value of type 'float' for argument 'a' but instead found type 'str'.
  
  aten::eq.int_float(int a, float b) -> (bool):
  Expected a value of type 'int' for argument 'a' but instead found type 'str'.
  
  aten::eq.float_int(float a, int b) -> (bool):
  Expected a value of type 'float' for argument 'a' but instead found type 'str'.
  
  aten::eq.float_complex(float a, complex b) -> (bool):
  Expected a value of type 'float' for argument 'a' but instead found type 'str'.
  
  aten::eq.complex_float(complex a, float b) -> (bool):
  Expected a value of type 'complex' for argument 'a' but instead found type 'str'.
  
  aten::eq(Scalar a, Scalar b) -> (bool):
  Expected a value of type 'number' for argument 'a' but instead found type 'str'.
  
  aten::eq.str(str a, str b) -> (bool):
  Expected a value of type 'str' for argument 'b' but instead found type 'Enum<__torch__.Color>'.
  
  aten::eq.float_list(float[] a, float[] b) -> (bool):
  Expected a value of type 'List[float]' for argument 'a' but instead found type 'str'.
  
  aten::eq.Tensor_list(Tensor[] a, Tensor[] b) -> (bool):
  Expected a value of type 'List[Tensor]' for argument 'a' but instead found type 'str'.
  
  aten::eq.bool_list(bool[] a, bool[] b) -> (bool):
  Expected a value of type 'List[bool]' for argument 'a' but instead found type 'str'.
  
  aten::eq.str_list(str[] a, str[] b) -> (bool):
  Expected a value of type 'List[str]' for argument 'a' but instead found type 'str'.
  
  eq(float a, Tensor b) -> (Tensor):
  Expected a value of type 'float' for argument 'a' but instead found type 'str'.
  
  eq(int a, Tensor b) -> (Tensor):
  Expected a value of type 'int' for argument 'a' but instead found type 'str'.

The original call is:
  File "torch_script_loss.py", line 71
    if x == Color.GREEN.value:
        return "green"
    if x == Color.RED:
       ~~~~~~~~~~~~~~ <--- HERE
        return "red"
```

### Versions

Collecting environment information...
PyTorch version: 1.12.0a0+2c916ef
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.3
Libc version: glibc-2.31

Python version: 3.8.12 | packaged by conda-forge | (default, Jan 30 2022, 23:42:07)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-30-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.6.112

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.942
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] pytorch-ignite==0.4.8
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.12.0a0+2c916ef
[pip3] torch-tensorrt==1.1.0a0
[pip3] torchinfo==1.6.5
[pip3] torchtext==0.12.0a0
[pip3] torchvision==0.13.0a0
[conda] magma-cuda110             2.5.2                         5    local
[conda] mkl                       2019.5                      281    conda-forge
[conda] mkl-include               2019.5                      281    conda-forge
[conda] numpy                     1.22.3           py38h05e7239_0    conda-forge
[conda] pytorch-ignite            0.4.8                    pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.12.0a0+2c916ef          pypi_0    pypi
[conda] torch-tensorrt            1.1.0a0                  pypi_0    pypi
[conda] torchinfo                 1.6.5                    pypi_0    pypi
[conda] torchtext                 0.12.0a0                 pypi_0    pypi
[conda] torchvision               0.13.0a0                 pypi_0    pypi