We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
New attribute `self._load_hook` in linear class since 2.03 will raise KeyError when executing `load_state_dict` fucntion
### 🐛 Describe the bug

In Pytorch, the function `load_state_dict(state_dict, strict)` allows empty dict `state_dict=={}` when `strict` is False. 
However, from version 2.03 the linear class in `torch_geometric.nn.dense.linear.py` has a new attribute `self._load_hook`, and when we execute `Linear(xxxx).load_state_dict({}, strict=False)`, the linear class will execute the `self._lazy_load_hook` function as follows
```
    def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,
                        missing_keys, unexpected_keys, error_msgs):

        weight = state_dict[prefix + 'weight']
        if is_uninitialized_parameter(weight):
            self.in_channels = -1
            self.weight = nn.parameter.UninitializedParameter()
            if not hasattr(self, '_hook'):
                self._hook = self.register_forward_pre_hook(
                    self.initialize_parameters)

        elif is_uninitialized_parameter(self.weight):
            self.in_channels = weight.size(-1)
            self.weight.materialize((self.out_channels, self.in_channels))
            if hasattr(self, '_hook'):
                self._hook.remove()
                delattr(self, '_hook')
```
Since the `state_dict` is empty, the line `weight = state_dict[prefix + 'weight']` will report KeyError. 

### Environment

* PyG version:
* PyTorch version:
* OS:
* Python version:
* CUDA/cuDNN version:
* How you installed PyTorch and PyG (`conda`, `pip`, source):
* Any other relevant information (*e.g.*, version of `torch-scatter`):


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch_geometric/nn/dense/linear.py`
Content:
```
1 import copy
2 import math
3 from typing import Any, Optional
4 
5 import torch
6 import torch.nn.functional as F
7 from torch import Tensor, nn
8 from torch.nn.parameter import Parameter
9 
10 from torch_geometric.nn import inits
11 
12 
13 def is_uninitialized_parameter(x: Any) -> bool:
14     if not hasattr(nn.parameter, 'UninitializedParameter'):
15         return False
16     return isinstance(x, nn.parameter.UninitializedParameter)
17 
18 
19 class Linear(torch.nn.Module):
20     r"""Applies a linear tranformation to the incoming data
21 
22     .. math::
23         \mathbf{x}^{\prime} = \mathbf{x} \mathbf{W}^{\top} + \mathbf{b}
24 
25     similar to :class:`torch.nn.Linear`.
26     It supports lazy initialization and customizable weight and bias
27     initialization.
28 
29     Args:
30         in_channels (int): Size of each input sample. Will be initialized
31             lazily in case it is given as :obj:`-1`.
32         out_channels (int): Size of each output sample.
33         bias (bool, optional): If set to :obj:`False`, the layer will not learn
34             an additive bias. (default: :obj:`True`)
35         weight_initializer (str, optional): The initializer for the weight
36             matrix (:obj:`"glorot"`, :obj:`"uniform"`, :obj:`"kaiming_uniform"`
37             or :obj:`None`).
38             If set to :obj:`None`, will match default weight initialization of
39             :class:`torch.nn.Linear`. (default: :obj:`None`)
40         bias_initializer (str, optional): The initializer for the bias vector
41             (:obj:`"zeros"` or :obj:`None`).
42             If set to :obj:`None`, will match default bias initialization of
43             :class:`torch.nn.Linear`. (default: :obj:`None`)
44 
45     Shapes:
46         - **input:** features :math:`(*, F_{in})`
47         - **output:** features :math:`(*, F_{out})`
48     """
49     def __init__(self, in_channels: int, out_channels: int, bias: bool = True,
50                  weight_initializer: Optional[str] = None,
51                  bias_initializer: Optional[str] = None):
52         super().__init__()
53         self.in_channels = in_channels
54         self.out_channels = out_channels
55         self.weight_initializer = weight_initializer
56         self.bias_initializer = bias_initializer
57 
58         if in_channels > 0:
59             self.weight = Parameter(torch.Tensor(out_channels, in_channels))
60         else:
61             self.weight = nn.parameter.UninitializedParameter()
62             self._hook = self.register_forward_pre_hook(
63                 self.initialize_parameters)
64 
65         if bias:
66             self.bias = Parameter(torch.Tensor(out_channels))
67         else:
68             self.register_parameter('bias', None)
69 
70         self._load_hook = self._register_load_state_dict_pre_hook(
71             self._lazy_load_hook)
72 
73         self.reset_parameters()
74 
75     def __deepcopy__(self, memo):
76         out = Linear(self.in_channels, self.out_channels, self.bias
77                      is not None, self.weight_initializer,
78                      self.bias_initializer)
79         if self.in_channels > 0:
80             out.weight = copy.deepcopy(self.weight, memo)
81         if self.bias is not None:
82             out.bias = copy.deepcopy(self.bias, memo)
83         return out
84 
85     def reset_parameters(self):
86         if self.in_channels <= 0:
87             pass
88         elif self.weight_initializer == 'glorot':
89             inits.glorot(self.weight)
90         elif self.weight_initializer == 'uniform':
91             bound = 1.0 / math.sqrt(self.weight.size(-1))
92             torch.nn.init.uniform_(self.weight.data, -bound, bound)
93         elif self.weight_initializer == 'kaiming_uniform':
94             inits.kaiming_uniform(self.weight, fan=self.in_channels,
95                                   a=math.sqrt(5))
96         elif self.weight_initializer is None:
97             inits.kaiming_uniform(self.weight, fan=self.in_channels,
98                                   a=math.sqrt(5))
99         else:
100             raise RuntimeError(f"Linear layer weight initializer "
101                                f"'{self.weight_initializer}' is not supported")
102 
103         if self.bias is None or self.in_channels <= 0:
104             pass
105         elif self.bias_initializer == 'zeros':
106             inits.zeros(self.bias)
107         elif self.bias_initializer is None:
108             inits.uniform(self.in_channels, self.bias)
109         else:
110             raise RuntimeError(f"Linear layer bias initializer "
111                                f"'{self.bias_initializer}' is not supported")
112 
113     def forward(self, x: Tensor) -> Tensor:
114         r"""
115         Args:
116             x (Tensor): The features.
117         """
118         return F.linear(x, self.weight, self.bias)
119 
120     @torch.no_grad()
121     def initialize_parameters(self, module, input):
122         if is_uninitialized_parameter(self.weight):
123             self.in_channels = input[0].size(-1)
124             self.weight.materialize((self.out_channels, self.in_channels))
125             self.reset_parameters()
126         self._hook.remove()
127         delattr(self, '_hook')
128 
129     def _save_to_state_dict(self, destination, prefix, keep_vars):
130         if is_uninitialized_parameter(self.weight):
131             destination[prefix + 'weight'] = self.weight
132         else:
133             destination[prefix + 'weight'] = self.weight.detach()
134         if self.bias is not None:
135             destination[prefix + 'bias'] = self.bias.detach()
136 
137     def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,
138                         missing_keys, unexpected_keys, error_msgs):
139 
140         weight = state_dict[prefix + 'weight']
141         if is_uninitialized_parameter(weight):
142             self.in_channels = -1
143             self.weight = nn.parameter.UninitializedParameter()
144             if not hasattr(self, '_hook'):
145                 self._hook = self.register_forward_pre_hook(
146                     self.initialize_parameters)
147 
148         elif is_uninitialized_parameter(self.weight):
149             self.in_channels = weight.size(-1)
150             self.weight.materialize((self.out_channels, self.in_channels))
151             if hasattr(self, '_hook'):
152                 self._hook.remove()
153                 delattr(self, '_hook')
154 
155     def __repr__(self) -> str:
156         return (f'{self.__class__.__name__}({self.in_channels}, '
157                 f'{self.out_channels}, bias={self.bias is not None})')
158 
159 
160 class HeteroLinear(torch.nn.Module):
161     r"""Applies separate linear tranformations to the incoming data according
162     to types
163 
164     .. math::
165         \mathbf{x}^{\prime}_{\kappa} = \mathbf{x}_{\kappa}
166         \mathbf{W}^{\top}_{\kappa} + \mathbf{b}_{\kappa}
167 
168     for type :math:`\kappa`.
169     It supports lazy initialization and customizable weight and bias
170     initialization.
171 
172     Args:
173         in_channels (int): Size of each input sample. Will be initialized
174             lazily in case it is given as :obj:`-1`.
175         out_channels (int): Size of each output sample.
176         num_types (int): The number of types.
177         **kwargs (optional): Additional arguments of
178             :class:`torch_geometric.nn.Linear`.
179 
180     Shapes:
181         - **input:**
182           features :math:`(*, F_{in})`,
183           type vector :math:`(*)`
184         - **output:** features :math:`(*, F_{out})`
185     """
186     def __init__(self, in_channels: int, out_channels: int, num_types: int,
187                  **kwargs):
188         super().__init__()
189 
190         self.in_channels = in_channels
191         self.out_channels = out_channels
192 
193         self.lins = torch.nn.ModuleList([
194             Linear(in_channels, out_channels, **kwargs)
195             for _ in range(num_types)
196         ])
197 
198         self.reset_parameters()
199 
200     def reset_parameters(self):
201         for lin in self.lins:
202             lin.reset_parameters()
203 
204     def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
205         r"""
206         Args:
207             x (Tensor): The input features.
208             type_vec (LongTensor): A vector that maps each entry to a type.
209         """
210         out = x.new_empty(x.size(0), self.out_channels)
211         for i, lin in enumerate(self.lins):
212             mask = type_vec == i
213             out[mask] = lin(x[mask])
214         return out
215 
216     def __repr__(self) -> str:
217         return (f'{self.__class__.__name__}({self.in_channels}, '
218                 f'{self.out_channels}, bias={self.lins[0].bias is not None})')
219
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -137,15 +137,16 @@
     def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,
                         missing_keys, unexpected_keys, error_msgs):
 
-        weight = state_dict[prefix + 'weight']
-        if is_uninitialized_parameter(weight):
+        weight = state_dict.get(prefix + 'weight', None)
+
+        if weight is not None and is_uninitialized_parameter(weight):
             self.in_channels = -1
             self.weight = nn.parameter.UninitializedParameter()
             if not hasattr(self, '_hook'):
                 self._hook = self.register_forward_pre_hook(
                     self.initialize_parameters)
 
-        elif is_uninitialized_parameter(self.weight):
+        elif weight is not None and is_uninitialized_parameter(self.weight):
             self.in_channels = weight.size(-1)
             self.weight.materialize((self.out_channels, self.in_channels))
             if hasattr(self, '_hook'):

{"golden_diff": "diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py\n--- a/torch_geometric/nn/dense/linear.py\n+++ b/torch_geometric/nn/dense/linear.py\n@@ -137,15 +137,16 @@\n     def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,\n                         missing_keys, unexpected_keys, error_msgs):\n \n-        weight = state_dict[prefix + 'weight']\n-        if is_uninitialized_parameter(weight):\n+        weight = state_dict.get(prefix + 'weight', None)\n+\n+        if weight is not None and is_uninitialized_parameter(weight):\n             self.in_channels = -1\n             self.weight = nn.parameter.UninitializedParameter()\n             if not hasattr(self, '_hook'):\n                 self._hook = self.register_forward_pre_hook(\n                     self.initialize_parameters)\n \n-        elif is_uninitialized_parameter(self.weight):\n+        elif weight is not None and is_uninitialized_parameter(self.weight):\n             self.in_channels = weight.size(-1)\n             self.weight.materialize((self.out_channels, self.in_channels))\n             if hasattr(self, '_hook'):\n", "issue": "New attribute `self._load_hook` in linear class since 2.03 will raise KeyError when executing `load_state_dict` fucntion\n### \ud83d\udc1b Describe the bug\r\n\r\nIn Pytorch, the function `load_state_dict(state_dict, strict)` allows empty dict `state_dict=={}` when `strict` is False. \r\nHowever, from version 2.03 the linear class in `torch_geometric.nn.dense.linear.py` has a new attribute `self._load_hook`, and when we execute `Linear(xxxx).load_state_dict({}, strict=False)`, the linear class will execute the `self._lazy_load_hook` function as follows\r\n```\r\n    def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,\r\n                        missing_keys, unexpected_keys, error_msgs):\r\n\r\n        weight = state_dict[prefix + 'weight']\r\n        if is_uninitialized_parameter(weight):\r\n            self.in_channels = -1\r\n            self.weight = nn.parameter.UninitializedParameter()\r\n            if not hasattr(self, '_hook'):\r\n                self._hook = self.register_forward_pre_hook(\r\n                    self.initialize_parameters)\r\n\r\n        elif is_uninitialized_parameter(self.weight):\r\n            self.in_channels = weight.size(-1)\r\n            self.weight.materialize((self.out_channels, self.in_channels))\r\n            if hasattr(self, '_hook'):\r\n                self._hook.remove()\r\n                delattr(self, '_hook')\r\n```\r\nSince the `state_dict` is empty, the line `weight = state_dict[prefix + 'weight']` will report KeyError. \r\n\r\n### Environment\r\n\r\n* PyG version:\r\n* PyTorch version:\r\n* OS:\r\n* Python version:\r\n* CUDA/cuDNN version:\r\n* How you installed PyTorch and PyG (`conda`, `pip`, source):\r\n* Any other relevant information (*e.g.*, version of `torch-scatter`):\r\n\n", "before_files": [{"content": "import copy\nimport math\nfrom typing import Any, Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import Tensor, nn\nfrom torch.nn.parameter import Parameter\n\nfrom torch_geometric.nn import inits\n\n\ndef is_uninitialized_parameter(x: Any) -> bool:\n    if not hasattr(nn.parameter, 'UninitializedParameter'):\n        return False\n    return isinstance(x, nn.parameter.UninitializedParameter)\n\n\nclass Linear(torch.nn.Module):\n    r\"\"\"Applies a linear tranformation to the incoming data\n\n    .. math::\n        \\mathbf{x}^{\\prime} = \\mathbf{x} \\mathbf{W}^{\\top} + \\mathbf{b}\n\n    similar to :class:`torch.nn.Linear`.\n    It supports lazy initialization and customizable weight and bias\n    initialization.\n\n    Args:\n        in_channels (int): Size of each input sample. Will be initialized\n            lazily in case it is given as :obj:`-1`.\n        out_channels (int): Size of each output sample.\n        bias (bool, optional): If set to :obj:`False`, the layer will not learn\n            an additive bias. (default: :obj:`True`)\n        weight_initializer (str, optional): The initializer for the weight\n            matrix (:obj:`\"glorot\"`, :obj:`\"uniform\"`, :obj:`\"kaiming_uniform\"`\n            or :obj:`None`).\n            If set to :obj:`None`, will match default weight initialization of\n            :class:`torch.nn.Linear`. (default: :obj:`None`)\n        bias_initializer (str, optional): The initializer for the bias vector\n            (:obj:`\"zeros\"` or :obj:`None`).\n            If set to :obj:`None`, will match default bias initialization of\n            :class:`torch.nn.Linear`. (default: :obj:`None`)\n\n    Shapes:\n        - **input:** features :math:`(*, F_{in})`\n        - **output:** features :math:`(*, F_{out})`\n    \"\"\"\n    def __init__(self, in_channels: int, out_channels: int, bias: bool = True,\n                 weight_initializer: Optional[str] = None,\n                 bias_initializer: Optional[str] = None):\n        super().__init__()\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.weight_initializer = weight_initializer\n        self.bias_initializer = bias_initializer\n\n        if in_channels > 0:\n            self.weight = Parameter(torch.Tensor(out_channels, in_channels))\n        else:\n            self.weight = nn.parameter.UninitializedParameter()\n            self._hook = self.register_forward_pre_hook(\n                self.initialize_parameters)\n\n        if bias:\n            self.bias = Parameter(torch.Tensor(out_channels))\n        else:\n            self.register_parameter('bias', None)\n\n        self._load_hook = self._register_load_state_dict_pre_hook(\n            self._lazy_load_hook)\n\n        self.reset_parameters()\n\n    def __deepcopy__(self, memo):\n        out = Linear(self.in_channels, self.out_channels, self.bias\n                     is not None, self.weight_initializer,\n                     self.bias_initializer)\n        if self.in_channels > 0:\n            out.weight = copy.deepcopy(self.weight, memo)\n        if self.bias is not None:\n            out.bias = copy.deepcopy(self.bias, memo)\n        return out\n\n    def reset_parameters(self):\n        if self.in_channels <= 0:\n            pass\n        elif self.weight_initializer == 'glorot':\n            inits.glorot(self.weight)\n        elif self.weight_initializer == 'uniform':\n            bound = 1.0 / math.sqrt(self.weight.size(-1))\n            torch.nn.init.uniform_(self.weight.data, -bound, bound)\n        elif self.weight_initializer == 'kaiming_uniform':\n            inits.kaiming_uniform(self.weight, fan=self.in_channels,\n                                  a=math.sqrt(5))\n        elif self.weight_initializer is None:\n            inits.kaiming_uniform(self.weight, fan=self.in_channels,\n                                  a=math.sqrt(5))\n        else:\n            raise RuntimeError(f\"Linear layer weight initializer \"\n                               f\"'{self.weight_initializer}' is not supported\")\n\n        if self.bias is None or self.in_channels <= 0:\n            pass\n        elif self.bias_initializer == 'zeros':\n            inits.zeros(self.bias)\n        elif self.bias_initializer is None:\n            inits.uniform(self.in_channels, self.bias)\n        else:\n            raise RuntimeError(f\"Linear layer bias initializer \"\n                               f\"'{self.bias_initializer}' is not supported\")\n\n    def forward(self, x: Tensor) -> Tensor:\n        r\"\"\"\n        Args:\n            x (Tensor): The features.\n        \"\"\"\n        return F.linear(x, self.weight, self.bias)\n\n    @torch.no_grad()\n    def initialize_parameters(self, module, input):\n        if is_uninitialized_parameter(self.weight):\n            self.in_channels = input[0].size(-1)\n            self.weight.materialize((self.out_channels, self.in_channels))\n            self.reset_parameters()\n        self._hook.remove()\n        delattr(self, '_hook')\n\n    def _save_to_state_dict(self, destination, prefix, keep_vars):\n        if is_uninitialized_parameter(self.weight):\n            destination[prefix + 'weight'] = self.weight\n        else:\n            destination[prefix + 'weight'] = self.weight.detach()\n        if self.bias is not None:\n            destination[prefix + 'bias'] = self.bias.detach()\n\n    def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,\n                        missing_keys, unexpected_keys, error_msgs):\n\n        weight = state_dict[prefix + 'weight']\n        if is_uninitialized_parameter(weight):\n            self.in_channels = -1\n            self.weight = nn.parameter.UninitializedParameter()\n            if not hasattr(self, '_hook'):\n                self._hook = self.register_forward_pre_hook(\n                    self.initialize_parameters)\n\n        elif is_uninitialized_parameter(self.weight):\n            self.in_channels = weight.size(-1)\n            self.weight.materialize((self.out_channels, self.in_channels))\n            if hasattr(self, '_hook'):\n                self._hook.remove()\n                delattr(self, '_hook')\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels}, bias={self.bias is not None})')\n\n\nclass HeteroLinear(torch.nn.Module):\n    r\"\"\"Applies separate linear tranformations to the incoming data according\n    to types\n\n    .. math::\n        \\mathbf{x}^{\\prime}_{\\kappa} = \\mathbf{x}_{\\kappa}\n        \\mathbf{W}^{\\top}_{\\kappa} + \\mathbf{b}_{\\kappa}\n\n    for type :math:`\\kappa`.\n    It supports lazy initialization and customizable weight and bias\n    initialization.\n\n    Args:\n        in_channels (int): Size of each input sample. Will be initialized\n            lazily in case it is given as :obj:`-1`.\n        out_channels (int): Size of each output sample.\n        num_types (int): The number of types.\n        **kwargs (optional): Additional arguments of\n            :class:`torch_geometric.nn.Linear`.\n\n    Shapes:\n        - **input:**\n          features :math:`(*, F_{in})`,\n          type vector :math:`(*)`\n        - **output:** features :math:`(*, F_{out})`\n    \"\"\"\n    def __init__(self, in_channels: int, out_channels: int, num_types: int,\n                 **kwargs):\n        super().__init__()\n\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n\n        self.lins = torch.nn.ModuleList([\n            Linear(in_channels, out_channels, **kwargs)\n            for _ in range(num_types)\n        ])\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        for lin in self.lins:\n            lin.reset_parameters()\n\n    def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:\n        r\"\"\"\n        Args:\n            x (Tensor): The input features.\n            type_vec (LongTensor): A vector that maps each entry to a type.\n        \"\"\"\n        out = x.new_empty(x.size(0), self.out_channels)\n        for i, lin in enumerate(self.lins):\n            mask = type_vec == i\n            out[mask] = lin(x[mask])\n        return out\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels}, bias={self.lins[0].bias is not None})')\n", "path": "torch_geometric/nn/dense/linear.py"}], "after_files": [{"content": "import copy\nimport math\nfrom typing import Any, Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import Tensor, nn\nfrom torch.nn.parameter import Parameter\n\nfrom torch_geometric.nn import inits\n\n\ndef is_uninitialized_parameter(x: Any) -> bool:\n    if not hasattr(nn.parameter, 'UninitializedParameter'):\n        return False\n    return isinstance(x, nn.parameter.UninitializedParameter)\n\n\nclass Linear(torch.nn.Module):\n    r\"\"\"Applies a linear tranformation to the incoming data\n\n    .. math::\n        \\mathbf{x}^{\\prime} = \\mathbf{x} \\mathbf{W}^{\\top} + \\mathbf{b}\n\n    similar to :class:`torch.nn.Linear`.\n    It supports lazy initialization and customizable weight and bias\n    initialization.\n\n    Args:\n        in_channels (int): Size of each input sample. Will be initialized\n            lazily in case it is given as :obj:`-1`.\n        out_channels (int): Size of each output sample.\n        bias (bool, optional): If set to :obj:`False`, the layer will not learn\n            an additive bias. (default: :obj:`True`)\n        weight_initializer (str, optional): The initializer for the weight\n            matrix (:obj:`\"glorot\"`, :obj:`\"uniform\"`, :obj:`\"kaiming_uniform\"`\n            or :obj:`None`).\n            If set to :obj:`None`, will match default weight initialization of\n            :class:`torch.nn.Linear`. (default: :obj:`None`)\n        bias_initializer (str, optional): The initializer for the bias vector\n            (:obj:`\"zeros\"` or :obj:`None`).\n            If set to :obj:`None`, will match default bias initialization of\n            :class:`torch.nn.Linear`. (default: :obj:`None`)\n\n    Shapes:\n        - **input:** features :math:`(*, F_{in})`\n        - **output:** features :math:`(*, F_{out})`\n    \"\"\"\n    def __init__(self, in_channels: int, out_channels: int, bias: bool = True,\n                 weight_initializer: Optional[str] = None,\n                 bias_initializer: Optional[str] = None):\n        super().__init__()\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.weight_initializer = weight_initializer\n        self.bias_initializer = bias_initializer\n\n        if in_channels > 0:\n            self.weight = Parameter(torch.Tensor(out_channels, in_channels))\n        else:\n            self.weight = nn.parameter.UninitializedParameter()\n            self._hook = self.register_forward_pre_hook(\n                self.initialize_parameters)\n\n        if bias:\n            self.bias = Parameter(torch.Tensor(out_channels))\n        else:\n            self.register_parameter('bias', None)\n\n        self._load_hook = self._register_load_state_dict_pre_hook(\n            self._lazy_load_hook)\n\n        self.reset_parameters()\n\n    def __deepcopy__(self, memo):\n        out = Linear(self.in_channels, self.out_channels, self.bias\n                     is not None, self.weight_initializer,\n                     self.bias_initializer)\n        if self.in_channels > 0:\n            out.weight = copy.deepcopy(self.weight, memo)\n        if self.bias is not None:\n            out.bias = copy.deepcopy(self.bias, memo)\n        return out\n\n    def reset_parameters(self):\n        if self.in_channels <= 0:\n            pass\n        elif self.weight_initializer == 'glorot':\n            inits.glorot(self.weight)\n        elif self.weight_initializer == 'uniform':\n            bound = 1.0 / math.sqrt(self.weight.size(-1))\n            torch.nn.init.uniform_(self.weight.data, -bound, bound)\n        elif self.weight_initializer == 'kaiming_uniform':\n            inits.kaiming_uniform(self.weight, fan=self.in_channels,\n                                  a=math.sqrt(5))\n        elif self.weight_initializer is None:\n            inits.kaiming_uniform(self.weight, fan=self.in_channels,\n                                  a=math.sqrt(5))\n        else:\n            raise RuntimeError(f\"Linear layer weight initializer \"\n                               f\"'{self.weight_initializer}' is not supported\")\n\n        if self.bias is None or self.in_channels <= 0:\n            pass\n        elif self.bias_initializer == 'zeros':\n            inits.zeros(self.bias)\n        elif self.bias_initializer is None:\n            inits.uniform(self.in_channels, self.bias)\n        else:\n            raise RuntimeError(f\"Linear layer bias initializer \"\n                               f\"'{self.bias_initializer}' is not supported\")\n\n    def forward(self, x: Tensor) -> Tensor:\n        r\"\"\"\n        Args:\n            x (Tensor): The features.\n        \"\"\"\n        return F.linear(x, self.weight, self.bias)\n\n    @torch.no_grad()\n    def initialize_parameters(self, module, input):\n        if is_uninitialized_parameter(self.weight):\n            self.in_channels = input[0].size(-1)\n            self.weight.materialize((self.out_channels, self.in_channels))\n            self.reset_parameters()\n        self._hook.remove()\n        delattr(self, '_hook')\n\n    def _save_to_state_dict(self, destination, prefix, keep_vars):\n        if is_uninitialized_parameter(self.weight):\n            destination[prefix + 'weight'] = self.weight\n        else:\n            destination[prefix + 'weight'] = self.weight.detach()\n        if self.bias is not None:\n            destination[prefix + 'bias'] = self.bias.detach()\n\n    def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,\n                        missing_keys, unexpected_keys, error_msgs):\n\n        weight = state_dict.get(prefix + 'weight', None)\n\n        if weight is not None and is_uninitialized_parameter(weight):\n            self.in_channels = -1\n            self.weight = nn.parameter.UninitializedParameter()\n            if not hasattr(self, '_hook'):\n                self._hook = self.register_forward_pre_hook(\n                    self.initialize_parameters)\n\n        elif weight is not None and is_uninitialized_parameter(self.weight):\n            self.in_channels = weight.size(-1)\n            self.weight.materialize((self.out_channels, self.in_channels))\n            if hasattr(self, '_hook'):\n                self._hook.remove()\n                delattr(self, '_hook')\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels}, bias={self.bias is not None})')\n\n\nclass HeteroLinear(torch.nn.Module):\n    r\"\"\"Applies separate linear tranformations to the incoming data according\n    to types\n\n    .. math::\n        \\mathbf{x}^{\\prime}_{\\kappa} = \\mathbf{x}_{\\kappa}\n        \\mathbf{W}^{\\top}_{\\kappa} + \\mathbf{b}_{\\kappa}\n\n    for type :math:`\\kappa`.\n    It supports lazy initialization and customizable weight and bias\n    initialization.\n\n    Args:\n        in_channels (int): Size of each input sample. Will be initialized\n            lazily in case it is given as :obj:`-1`.\n        out_channels (int): Size of each output sample.\n        num_types (int): The number of types.\n        **kwargs (optional): Additional arguments of\n            :class:`torch_geometric.nn.Linear`.\n\n    Shapes:\n        - **input:**\n          features :math:`(*, F_{in})`,\n          type vector :math:`(*)`\n        - **output:** features :math:`(*, F_{out})`\n    \"\"\"\n    def __init__(self, in_channels: int, out_channels: int, num_types: int,\n                 **kwargs):\n        super().__init__()\n\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n\n        self.lins = torch.nn.ModuleList([\n            Linear(in_channels, out_channels, **kwargs)\n            for _ in range(num_types)\n        ])\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        for lin in self.lins:\n            lin.reset_parameters()\n\n    def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:\n        r\"\"\"\n        Args:\n            x (Tensor): The input features.\n            type_vec (LongTensor): A vector that maps each entry to a type.\n        \"\"\"\n        out = x.new_empty(x.size(0), self.out_channels)\n        for i, lin in enumerate(self.lins):\n            mask = type_vec == i\n            out[mask] = lin(x[mask])\n        return out\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels}, bias={self.lins[0].bias is not None})')\n", "path": "torch_geometric/nn/dense/linear.py"}]}