class_docstr: Applies separate linear tranformations to the incoming data according
  to types.
  
  For type :math:`\kappa`, it computes
  
  .. math::
      \mathbf{x}^{\prime}_{\kappa} = \mathbf{x}_{\kappa}
      \mathbf{W}^{\top}_{\kappa} + \mathbf{b}_{\kappa}.
  
  It supports lazy initialization and customizable weight and bias
  initialization.
  
  Args:
      in_channels (int): Size of each input sample. Will be initialized
          lazily in case it is given as :obj:`-1`.
      out_channels (int): Size of each output sample.
      num_types (int): The number of types.
      is_sorted (bool, optional): If set to :obj:`True`, assumes that
          :obj:`type_vec` is sorted. This avoids internal re-sorting of the
          data and can improve runtime and memory efficiency.
          (default: :obj:`False`)
      **kwargs (optional): Additional arguments of
          :class:`torch_geometric.nn.Linear`.
  
  Shapes:
      - **input:**
        features :math:`(*, F_{in})`,
        type vector :math:`(*)`
      - **output:** features :math:`(*, F_{out})`
class_docstr_tok: ["Applies","separate","linear","tranformations","to","the","incoming","data","according","to","types",".","For","type",":","math",":","`","\\kappa","`",",","it","computes","..","math",":",":","\\mathbf","{","x","}","^","{","\\prime","}","_","{","\\kappa","}","=","\\mathbf","{","x","}","_","{","\\kappa","}","\\mathbf","{","W","}","^","{","\\top","}","_","{","\\kappa","}","+","\\mathbf","{","b","}","_","{","\\kappa","}",".","It","supports","lazy","initialization","and","customizable","weight","and","bias","initialization",".","Args",":","in_channels","(","int",")",":","Size","of","each","input","sample",".","Will","be","initialized","lazily","in","case","it","is","given","as",":","obj",":","`","-1","`",".","out_channels","(","int",")",":","Size","of","each","output","sample",".","num_types","(","int",")",":","The","number","of","types",".","is_sorted","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","assumes","that",":","obj",":","`","type_vec","`","is","sorted",".","This","avoids","internal","re-sorting","of","the","data","and","can","improve","runtime","and","memory","efficiency",".","(","default",":",":","obj",":","`","False","`",")","*","*","kwargs","(","optional",")",":","Additional","arguments","of",":","class",":","`","torch_geometric.nn.Linear","`",".","Shapes",":","-","*","*","input",":","*","*","features",":","math",":","`","(","*",",","F_","{","in","}",")","`",",","type","vector",":","math",":","`","(","*",")","`","-","*","*","output",":","*","*","features",":","math",":","`","(","*",",","F_","{","out","}",")","`"]
docstr: Forward pass.
  
  Args:
      x (torch.Tensor): The input features.
      type_vec (torch.Tensor): A vector that maps each entry to a type.
docstr_tok: ["Forward","pass",".","Args",":","x","(","torch.Tensor",")",":","The","input","features",".","type_vec","(","torch.Tensor",")",":","A","vector","that","maps","each","entry","to","a","type","."]
code: def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
      r"""Forward pass.
  
      Args:
          x (torch.Tensor): The input features.
          type_vec (torch.Tensor): A vector that maps each entry to a type.
      """
      use_segment_matmul = torch_geometric.backend.use_segment_matmul
      # If `use_segment_matmul` is not specified, use a simple heuristic to
      # determine whether `segment_matmul` can speed up computation given the
      # observed input sizes:
      if use_segment_matmul is None:
          if self._use_segment_matmul_heuristic_output is None:
              segment_count = scatter(
                  torch.ones_like(type_vec),
                  type_vec,
                  dim_size=self.num_types,
                  reduce="sum",
              )
  
              self._use_segment_matmul_heuristic_output = (
                  torch_geometric.backend.use_segment_matmul_heuristic(
                      num_segments=self.num_types,
                      max_segment_size=int(segment_count.max()),
                      in_channels=self.weight.size(1),
                      out_channels=self.weight.size(2),
                  )
              )
  
          assert self._use_segment_matmul_heuristic_output is not None
          use_segment_matmul = self._use_segment_matmul_heuristic_output
  
      if (
          use_segment_matmul
          and torch_geometric.typing.WITH_SEGMM
          and not is_compiling()
      ):
          assert self.weight is not None
  
          perm: Optional[Tensor] = None
          if not self.is_sorted:
              if (type_vec[1:] < type_vec[:-1]).any():
                  type_vec, perm = index_sort(type_vec, self.num_types)
                  x = x[perm]
  
          type_vec_ptr = index2ptr(type_vec, self.num_types)
          out = pyg_lib.ops.segment_matmul(x, type_vec_ptr, self.weight)
          if self.bias is not None:
              out += self.bias[type_vec]
  
          if perm is not None:  # Restore original order (if necessary).
              out_unsorted = torch.empty_like(out)
              out_unsorted[perm] = out
              out = out_unsorted
      else:
          out = x.new_empty(x.size(0), self.out_channels)
          for i in range(self.num_types):
              mask = type_vec == i
              if mask.numel() == 0:
                  continue
              subset_out = F.linear(x[mask], self.weight[i].T)
              # The data type may have changed with mixed precision:
              out[mask] = subset_out.to(out.dtype)
  
          if self.bias is not None:
              out += self.bias[type_vec]
      return out
  
imports: import copy
  import math
  from typing import Any, Dict, Optional, Union
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch.nn.parameter import Parameter
  import torch_geometric.backend
  import torch_geometric.typing
  from torch_geometric import is_compiling
  from torch_geometric.nn import inits
  from torch_geometric.typing import pyg_lib
  from torch_geometric.utils import index_sort, scatter
  from torch_geometric.utils.sparse import index2ptr
label_desc: Use image node_id 3 for calling the HeteroLinear obj's underlying member method code with example usage: obj.forward(x, type_vec) and returns: out