Repository: pytorch
Issue: `torch.jit.load` fails when function parameters use non-ASCII characters
Labels: oncall: jit

### 🐛 Describe the bug

When you want to load a JIT-compiled ScriptModule that has some function (in the recursive ScriptModule chain) which contains a parameter with a non-ASCII name, you will get a `RuntimeError` with no error description.

See the example below:
```python
import io

import torch
import torch.nn as nn


class MLP(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, output_dim)
    
    def forward(self, x_θ: torch.Tensor) -> torch.Tensor:
        x_θ = x_θ
        x_θ = torch.relu(self.fc1(x_θ))
        x_θ = torch.relu(self.fc2(x_θ))
        x_θ = self.fc3(x_θ)
        return x_θ

model = MLP(input_dim=2, output_dim=1)
jit_model = torch.jit.script(model)
buffer = jit_model.save_to_buffer()
reloaded_model = torch.jit.load(io.BytesIO(buffer))
```
The code above will produce the following error:
```
RuntimeError                              Traceback (most recent call last)
/private/home/nimashoghi/ocp-pytorch-lightning/basic-test.ipynb Cell 1' in <module>
     22 jit_model = torch.jit.script(model)
     23 buffer = jit_model.save_to_buffer()
---> 24 reloaded_model = torch.jit.load(io.BytesIO(buffer))

File ~/.conda/envs/ocp-pytorch-lightning/lib/python3.9/site-packages/torch/jit/_serialization.py:163, in load(f, map_location, _extra_files)
    161     cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files)
    162 else:
--> 163     cpp_module = torch._C.import_ir_module_from_buffer(
    164         cu, f.read(), map_location, _extra_files
    165     )
    167 # TODO: Pretty sure this approach loses ConstSequential status and such
    168 return wrap_cpp_module(cpp_module)

RuntimeError: 
```

gdb stacktrace:
```
#1  0x00007fffd2531a9d in torch::jit::Lexer::expected (this=0x555558047910, what=..., t=...) at ../torch/csrc/jit/frontend/lexer.h:462
#2  0x00007fffd25321f8 in torch::jit::Lexer::lexRaw (this=0x555558047910, whitespace_token=false) at ../torch/csrc/jit/frontend/lexer.h:552
#3  0x00007fffd2531c05 in torch::jit::Lexer::lex (this=0x555558047910) at ../torch/csrc/jit/frontend/lexer.h:487
#4  0x00007fffd2531759 in torch::jit::Lexer::next (this=0x555558047910) at ../torch/csrc/jit/frontend/lexer.h:436
#5  0x00007fffd2531b81 in torch::jit::Lexer::expect (this=0x555558047910, kind=266) at ../torch/csrc/jit/frontend/lexer.h:473
#6  0x00007fffd5d6be41 in torch::jit::ParserImpl::parseIdent (this=0x555558047910) at ../torch/csrc/jit/frontend/parser.cpp:53
#7  0x00007fffd5d6f990 in torch::jit::ParserImpl::parseFormalParam (this=0x555558047910, kwarg_only=false) at ../torch/csrc/jit/frontend/parser.cpp:451
#8  0x00007fffd5d72951 in torch::jit::ParserImpl::parseFormalParams()::{lambda()#1}::operator()() const (__closure=0x55555803f540)
    at ../torch/csrc/jit/frontend/parser.cpp:721
#9  0x00007fffd5d75a4b in std::_Function_handler<void (), torch::jit::ParserImpl::parseFormalParams()::{lambda()#1}>::_M_invoke(std::_Any_data const&) (
    __functor=...) at /usr/include/c++/8/bits/std_function.h:297
#10 0x00007fffe2126482 in std::function<void ()>::operator()() const (this=0x7fffffff9cd0) at /usr/include/c++/8/bits/std_function.h:687
#11 0x00007fffd5d6ea0a in torch::jit::ParserImpl::parseSequence(int, int, int, std::function<void ()> const&) (this=0x555558047910, begin=40, sep=44,
    end=41, parse=...) at ../torch/csrc/jit/frontend/parser.cpp:340
#12 0x00007fffd5d72a48 in torch::jit::ParserImpl::parseFormalParams (this=0x555558047910) at ../torch/csrc/jit/frontend/parser.cpp:717
#13 0x00007fffd5d72b36 in torch::jit::ParserImpl::parseDecl (this=0x555558047910) at ../torch/csrc/jit/frontend/parser.cpp:728
#14 0x00007fffd5d72fdb in torch::jit::ParserImpl::parseFunction (this=0x555558047910, is_method=true) at ../torch/csrc/jit/frontend/parser.cpp:756
#15 0x00007fffd5d714b6 in torch::jit::ParserImpl::parseStmt (this=0x555558047910, in_class=true) at ../torch/csrc/jit/frontend/parser.cpp:600
#16 0x00007fffd5d72703 in torch::jit::ParserImpl::parseStatements (this=0x555558047910, expect_indent=true, in_class=true)
    at ../torch/csrc/jit/frontend/parser.cpp:698
#17 0x00007fffd5d72dd3 in torch::jit::ParserImpl::parseClass (this=0x555558047910) at ../torch/csrc/jit/frontend/parser.cpp:748
#18 0x00007fffd5d697e8 in torch::jit::Parser::parseClass (this=0x7fffffffa898) at ../torch/csrc/jit/frontend/parser.cpp:813
#19 0x00007fffd62b4974 in torch::jit::SourceImporterImpl::parseSourceIfNeeded (this=0x555557fc3bd0, qualifier=...)
    at ../torch/csrc/jit/serialization/import_source.cpp:182
#20 0x00007fffd62b44b4 in torch::jit::SourceImporterImpl::findNamedType (this=0x555557fc3bd0, name=...)
    at ../torch/csrc/jit/serialization/import_source.cpp:135
#21 0x00007fffd62b5b3d in torch::jit::SourceImporterImpl::resolveType (this=0x555557fc3bd0, name=..., loc=...)
    at ../torch/csrc/jit/serialization/import_source.cpp:261
#22 0x00007fffd5d86c25 in torch::jit::ScriptTypeParser::parseTypeFromExpr (this=0x7fffffffb090, expr=...)
    at ../torch/csrc/jit/frontend/script_type_parser.cpp:230
#23 0x00007fffd5d87af9 in torch::jit::ScriptTypeParser::parseType (this=0x7fffffffb090, str=...) at ../torch/csrc/jit/frontend/script_type_parser.cpp:304
#24 0x00007fffd62bc6af in torch::jit::SourceImporter::loadType (this=0x7fffffffc660, name=...) at ../torch/csrc/jit/serialization/import_source.cpp:765
````
^ in other words, it's in the jit parser/lexer

This took me quite some time to debug, as (1) there's no error description and (2) torch.jit.save doesn't produce any errors. I can verify that it's only the parameter names that cause this (in other words, if I use non-ASCII variable names in the code but ASCII parameter names, that's OK).

### Versions

```
PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.10 | packaged by conda-forge | (main, Feb  1 2022, 21:24:11)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-81-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] functorch==0.0.1a0+8ffa328
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.5
[pip3] pytorch-lightning==1.6.0.dev0
[pip3] torch==1.10.0
[pip3] torch-cluster==1.5.9
[pip3] torch-geometric==2.0.3
[pip3] torch-scatter==2.0.9
[pip3] torch-sparse==0.6.12
[pip3] torch-spline-conv==1.2.1
[pip3] torchmetrics==0.7.2
[conda] blas                      2.113                       mkl    conda-forge
[conda] blas-devel                3.9.0            13_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1              ha36c431_10    conda-forge
[conda] libblas                   3.9.0            13_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            13_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            13_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            13_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h8d4b97c_803    conda-forge
[conda] mkl-devel                 2022.0.1           ha770c72_804    conda-forge
[conda] mkl-include               2022.0.1           h8d4b97c_803    conda-forge
[conda] mypy_extensions           0.4.3            py39hf3d152e_4    conda-forge
[conda] numpy                     1.21.5           py39haac66dc_0    conda-forge
[conda] pyg                       2.0.3           py39_torch_1.10.0_cu113    pyg
[conda] pytorch                   1.10.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-cluster           1.5.9           py39_torch_1.10.0_cu113    pyg
[conda] pytorch-lightning         1.6.0.dev0               pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-scatter           2.0.9           py39_torch_1.10.0_cu113    pyg
[conda] pytorch-sparse            0.6.12          py39_torch_1.10.0_cu113    pyg
[conda] pytorch-spline-conv       1.2.1           py39_torch_1.10.0_cu113    pyg
[conda] torchmetrics              0.7.2                    pypi_0    pypi
```
cc @davidberard98