Repository: pytorch
Issue: using Union[str, Tensor] as an argument to a torch.jit.script function
Labels: oncall: jit,triaged

### 🐛 Describe the bug

I'd like to be able to pass a `Dict[str, Union[str, Tensor]]` as an argument to a torch.jit.script function:

```py
from typing import Dict, Union
import torch
from torch import nn 


class GainAndAbs(nn.Module):
  
  def forward(self, x: torch.Tensor, params: Dict[str, Union[str, torch.Tensor]]):
    assert isinstance(params["mode"], str)
    assert isinstance(params["gain"], torch.Tensor)

    if params["mode"] == "gain":
      return x * params["gain"]
    
    elif params["mode"] == "abs":
        return x.abs()

model = GainAndAbs()
scripted_model = torch.jit.script(model)

params = {
    "mode": "gain",
    "gain": 2.0,
}

print(f"without torch.jit.script")
print(model(torch.tensor([1.0]), params))

print(f"with torch.jit.script")
print(scripted_model(torch.tensor([1.0]), params))
```

I am having trouble with the compiler resolving tensor values: 

```
Traceback (most recent call last):
  File "/Users/hugo/projects/pytensorjuce/ctrltest.py", line 19, in <module>
    scripted_model = torch.jit.script(model)
  File "/Users/hugo/miniconda3/lib/python3.9/site-packages/torch/jit/_script.py", line 1284, in script
    return torch.jit._recursive.create_script_module(
  File "/Users/hugo/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 480, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/Users/hugo/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 546, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/Users/hugo/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 397, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
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 'Union[Tensor, str]'.
  
  aten::eq.Scalar(Tensor self, Scalar other) -> Tensor:
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'Union[Tensor, 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 'Union[Tensor, 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 'Union[Tensor, str]'.
  
  aten::eq.int_list(int[] a, int[] b) -> bool:
  Expected a value of type 'List[int]' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.device(Device a, Device b) -> bool:
  Expected a value of type 'Device' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.bool(bool a, bool b) -> bool:
  Expected a value of type 'bool' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.enum(AnyEnumType a, AnyEnumType b) -> bool:
  Expected a value of type 'AnyEnumType' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.int(int a, int b) -> bool:
  Expected a value of type 'int' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.complex(complex a, complex b) -> bool:
  Expected a value of type 'complex' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.float(float a, float b) -> bool:
  Expected a value of type 'float' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.int_float(int a, float b) -> bool:
  Expected a value of type 'int' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.float_int(float a, int b) -> bool:
  Expected a value of type 'float' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.float_complex(float a, complex b) -> bool:
  Expected a value of type 'float' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.complex_float(complex a, float b) -> bool:
  Expected a value of type 'complex' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq(Scalar a, Scalar b) -> bool:
  Expected a value of type 'number' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.str(str a, str b) -> bool:
  Expected a value of type 'str' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.float_list(float[] a, float[] b) -> bool:
  Expected a value of type 'List[float]' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.Tensor_list(Tensor[] a, Tensor[] b) -> bool:
  Expected a value of type 'List[Tensor]' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.bool_list(bool[] a, bool[] b) -> bool:
  Expected a value of type 'List[bool]' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  aten::eq.str_list(str[] a, str[] b) -> bool:
  Expected a value of type 'List[str]' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  eq(float a, Tensor b) -> Tensor:
  Expected a value of type 'float' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  eq(int a, Tensor b) -> Tensor:
  Expected a value of type 'int' for argument 'a' but instead found type 'Union[Tensor, str]'.
  
  eq(complex a, Tensor b) -> Tensor:
  Expected a value of type 'complex' for argument 'a' but instead found type 'Union[Tensor, str]'.

The original call is:
  File "/Users/hugo/projects/pytensorjuce/ctrltest.py", line 12
    assert isinstance(params["gain"], torch.Tensor)
  
    if params["mode"] == "gain":
       ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
      return x * params["gain"]
```

I've tried adding assertions (as mentioned [here](https://github.com/pytorch/pytorch/issues/69434#issuecomment-1064306579)), but that didn't work. 


Any help with a workaround // solution would be appreciated! thanks!

### Versions


Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.0 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.24.2
Libc version: N/A

Python version: 3.9.12 (main, Apr  5 2022, 01:52:34)  [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-13.0-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1

Versions of relevant libraries:
[pip3] numpy==1.23.3
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[conda] numpy                     1.23.3                   pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torchaudio                2.0.2                    pypi_0    pypi

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel