class_docstr: Base class for creating message passing layers of the form
  
  .. math::
      \mathbf{x}_i^{\prime} = \gamma_{\mathbf{\Theta}} \left( \mathbf{x}_i,
      \square_{j \in \mathcal{N}(i)} \, \phi_{\mathbf{\Theta}}
      \left(\mathbf{x}_i, \mathbf{x}_j,\mathbf{e}_{j,i}\right) \right),
  
  where :math:`\square` denotes a differentiable, permutation invariant
  function, *e.g.*, sum, mean or max, and :math:`\gamma_{\mathbf{\Theta}}`
  and :math:`\phi_{\mathbf{\Theta}}` denote differentiable functions such as
  MLPs.
  See `here <https://pytorch-geometric.readthedocs.io/en/latest/notes/
  create_gnn.html>`__ for the accompanying tutorial.
  
  Args:
      aggr (string, optional): The aggregation scheme to use
          (:obj:`"add"`, :obj:`"mean"`, :obj:`"max"` or :obj:`None`).
          (default: :obj:`"add"`)
      flow (string, optional): The flow direction of message passing
          (:obj:`"source_to_target"` or :obj:`"target_to_source"`).
          (default: :obj:`"source_to_target"`)
      node_dim (int, optional): The axis along which to propagate.
          (default: :obj:`-2`)
docstr: Analyzes the :class:`MessagePassing` instance and produces a new
  jittable module.
  
  Args:
      typing (string, optional): If given, will generate a concrete
          instance with :meth:`forward` types based on :obj:`typing`,
          *e.g.*: :obj:`"(Tensor, Optional[Tensor]) -> Tensor"`.
docstr_tok: ["Analyzes","the",":","class",":","`","MessagePassing","`","instance","and","produces","a","new","jittable","module",".","Args",":","typing","(","string",",","optional",")",":","If","given",",","will","generate","a","concrete","instance","with",":","meth",":","`","forward","`","types","based","on",":","obj",":","`","typing","`",",","*","e.g",".","*",":",":","obj",":","`","``","(","Tensor",",","Optional","[","Tensor","]",")","-",">","Tensor","''","`","."]
code: def jittable(self, typing: Optional[str] = None):
      r"""Analyzes the :class:`MessagePassing` instance and produces a new
      jittable module.
  
      Args:
          typing (string, optional): If given, will generate a concrete
              instance with :meth:`forward` types based on :obj:`typing`,
              *e.g.*: :obj:`"(Tensor, Optional[Tensor]) -> Tensor"`.
      """
      # Find and parse `propagate()` types to format `{arg1: type1, ...}`.
      if hasattr(self, "propagate_type"):
          prop_types = {
              k: sanitize(str(v))
              for k, v in self.propagate_type.items()
          }
      else:
          source = inspect.getsource(self.__class__)
          match = re.search(r"#\s*propagate_type:\s*\((.*)\)", source)
          if match is None:
              raise TypeError(
                  "TorchScript support requires the definition of the types "
                  "passed to `propagate()`. Please specificy them via\n\n"
                  'propagate_type = {"arg1": type1, "arg2": type2, ... }\n\n'
                  "or via\n\n"
                  "# propagate_type: (arg1: type1, arg2: type2, ...)\n\n"
                  "inside the `MessagePassing` module."
              )
          prop_types = split_types_repr(match.group(1))
          prop_types = dict(
              [re.split(r"\s*:\s*", t) for t in prop_types]
          )
  
      # Parse `__collect__()` types to format `{arg:1, type1, ...}`.
      collect_types = self.inspector.types(
          ["message", "aggregate", "update"]
      )
  
      # Collect `forward()` header, body and @overload types.
      forward_types = parse_types(self.forward)
      forward_types = [resolve_types(*types) for types in forward_types]
      forward_types = list(chain.from_iterable(forward_types))
  
      keep_annotation = len(forward_types) < 2
      forward_header = func_header_repr(self.forward, keep_annotation)
      forward_body = func_body_repr(self.forward, keep_annotation)
  
      if keep_annotation:
          forward_types = []
      elif typing is not None:
          forward_types = []
          forward_body = 8 * " " + f"# type: {typing}\n{forward_body}"
  
      root = os.path.dirname(osp.realpath(__file__))
      with open(osp.join(root, "message_passing.jinja"), "r") as f:
          template = Template(f.read())
  
      uid = uuid1().hex[:6]
      cls_name = f"{self.__class__.__name__}Jittable_{uid}"
      jit_module_repr = template.render(
          uid=uid,
          module=str(self.__class__.__module__),
          cls_name=cls_name,
          parent_cls_name=self.__class__.__name__,
          prop_types=prop_types,
          collect_types=collect_types,
          user_args=self.__user_args__,
          forward_header=forward_header,
          forward_types=forward_types,
          forward_body=forward_body,
          msg_args=self.inspector.keys(["message"]),
          aggr_args=self.inspector.keys(["aggregate"]),
          msg_and_aggr_args=self.inspector.keys(
              ["message_and_aggregate"]
          ),
          update_args=self.inspector.keys(["update"]),
          check_input=inspect.getsource(self.__check_input__)[:-1],
          lift=inspect.getsource(self.__lift__)[:-1],
      )
  
      # Instantiate a class from the rendered JIT module representation.
      cls = class_from_module_repr(cls_name, jit_module_repr)
      module = cls.__new__(cls)
      module.__dict__ = self.__dict__.copy()
      module.jittable = None
  
      return module
  
imports: import os
  import re
  import inspect
  import os.path
  from uuid import uuid1
  from itertools import chain
  from inspect import Parameter
  from typing import List, Optional, Set
  from torch_geometric.typing import Adj, Size
  import torch
  from torch import Tensor
  from jinja2 import Template
  from torch_sparse import SparseTensor
  from torch_scatter import gather_csr, scatter, segment_csr
  from .utils.helpers import expand_left
  from .utils.jit import class_from_module_repr
  from .utils.typing import sanitize, split_types_repr, parse_types, resolve_types
  from .utils.inspector import Inspector, func_header_repr, func_body_repr
label_desc: Use image node_id 11 for calling the MessagePassing obj's underlying member method code with example usage: obj.jittable(typing) and returns: module