Repository: pyg-team/pytorch_geometric
Pull request #8461 (2023-11-28 09:55:43)

Add `EdgeIndex` tests for `torch.compile` and `torch.jit.script`; Add `is_torch_instance`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8461?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`d92ca41`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/d92ca41a5a2425f5c1afc700f443d5d3c0b107df?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.30% compared to head [(`417e786`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8461?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.68%.

> :exclamation: Current head 417e786 differs from pull request most recent head 274e596. Consider uploading reports for the commit 274e596 to get more accurate results


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8461      +/-   ##
==========================================
- Coverage   88.30%   87.68%   -0.63%     
==========================================
  Files         477      478       +1     
  Lines       29264    29266       +2     
==========================================
- Hits        25842    25661     -181     
- Misses       3422     3605     +183     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8461?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md
index c6396f138..96b37c770 100644
--- a/.github/CONTRIBUTING.md
+++ b/.github/CONTRIBUTING.md
@@ -4,7 +4,7 @@ If you are interested in contributing to PyG, your contributions will likely fal
 
 1. You want to implement a new feature:
    - In general, we accept any features as long as they fit the scope of this package. If you are unsure about this or need help on the design/implementation of your feature, post about it in an issue.
-2. You want to fix a bug:
+1. You want to fix a bug:
    - Feel free to send a Pull Request (PR) any time you encounter a bug. Please provide a clear and concise description of what the bug was. If you are unsure about if this is a bug at all or how to fix, post about it in an issue.
 
 Once you finish implementing a feature or bug-fix, please send a PR to https://github.com/pyg-team/pytorch_geometric.
@@ -23,7 +23,7 @@ To develop PyG on your machine, here are some tips:
    print(torch.__version__)
    ```
 
-2. *(Optional)* Follow the [installation instructions](https://github.com/pyg-team/pytorch_geometric#installation) to install `pyg-lib`, `torch-scatter`, `torch-sparse`, `torch-cluster` and `torch-spline-conv` (if you haven't already).
+1. *(Optional)* Follow the [installation instructions](https://github.com/pyg-team/pytorch_geometric#installation) to install `pyg-lib`, `torch-scatter`, `torch-sparse`, `torch-cluster` and `torch-spline-conv` (if you haven't already).
    Note that this step is optional and only necessary if you develop a feature that uses one of these libraries.
 
    ```bash
@@ -32,7 +32,7 @@ To develop PyG on your machine, here are some tips:
 
    where `${TORCH}` should be replaced by your PyTorch version (*e.g.*, `2.0.0`), and `${CUDA}` should be replaced by your CUDA version (*e.g.*, `cpu` or `cu118`).
 
-3. Uninstall all existing PyG installations.
+1. Uninstall all existing PyG installations.
    It is advised to run this command repeatedly to confirm that installations across all locations are properly removed.
 
    ```bash
@@ -40,20 +40,20 @@ To develop PyG on your machine, here are some tips:
    pip uninstall torch-geometric  # run this command twice
    ```
 
-4. Fork and clone the PyG repository:
+1. Fork and clone the PyG repository:
 
    ```bash
    git clone https://github.com/<your_username>/pytorch_geometric
    cd pytorch_geometric
    ```
 
-5. If you already cloned PyG from source, update it:
+1. If you already cloned PyG from source, update it:
 
    ```bash
    git pull
    ```
 
-6. Install PyG in editable mode:
+1. Install PyG in editable mode:
 
    ```bash
    pip install -e ".[dev,full]"
@@ -62,7 +62,7 @@ To develop PyG on your machine, here are some tips:
    This mode will symlink the Python files from the current local source tree into the Python install.
    Hence, if you modify a Python file, you do not need to re-install PyG again.
 
-7. Ensure that you have a working PyG installation by running the entire test suite with
+1. Ensure that you have a working PyG installation by running the entire test suite with
 
    ```bash
    pytest
@@ -70,7 +70,7 @@ To develop PyG on your machine, here are some tips:
 
    In case an error occurs, please first check if all sub-packages ([`pyg-lib`](https://github.com/pyg-team/pyg-lib), [`torch-scatter`](https://github.com/rusty1s/pytorch_scatter), [`torch-sparse`](https://github.com/rusty1s/pytorch_sparse), [`torch-cluster`](https://github.com/rusty1s/pytorch_cluster) and [`torch-spline-conv`](https://github.com/rusty1s/pytorch_spline_conv)) are on its latest reported version.
 
-8. Install pre-commit hooks:
+1. Install pre-commit hooks:
 
    ```bash
     pre-commit install
@@ -102,7 +102,7 @@ Everytime you send a Pull Request, your commit will be built and checked against
 
    If you do not want to format your code manually, we recommend to use [`yapf`](https://github.com/google/yapf).
 
-2. Ensure that the entire test suite passes and that code coverage roughly stays the same.
+1. Ensure that the entire test suite passes and that code coverage roughly stays the same.
    Please feel encouraged to provide a test with your submitted code.
    To test, either run
 
@@ -118,7 +118,7 @@ Everytime you send a Pull Request, your commit will be built and checked against
 
    (which runs a set of additional but time-consuming tests) dependening on your needs.
 
-3. Add your feature/bugfix to the [`CHANGELOG.md`](https://github.com/pyg-team/pytorch_geometric/blob/master/CHANGELOG.md?plain=1).
+1. Add your feature/bugfix to the [`CHANGELOG.md`](https://github.com/pyg-team/pytorch_geometric/blob/master/CHANGELOG.md?plain=1).
    If multiple PRs move towards integrating a single feature, it is advised to group them together into one bullet point.
 
 ## Building Documentation
@@ -126,11 +126,11 @@ Everytime you send a Pull Request, your commit will be built and checked against
 To build the documentation:
 
 1. [Build and install](#developing-pyg) PyG from source.
-2. Install [Sphinx](https://www.sphinx-doc.org/en/master/) theme via
+1. Install [Sphinx](https://www.sphinx-doc.org/en/master/) theme via
    ```bash
    pip install git+https://github.com/pyg-team/pyg_sphinx_theme.git
    ```
-3. Generate the documentation via:
+1. Generate the documentation via:
    ```bash
    cd docs
    make html
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b3aa7e3f3..a896986dd 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -63,3 +63,18 @@ repos:
       - id: ruff
         name: Ruff formatting
         args: [--fix, --exit-non-zero-on-fix]
+
+  - repo: https://github.com/executablebooks/mdformat
+    rev: 0.7.17
+    hooks:
+      - id: mdformat
+        name: Format Markdown
+        additional_dependencies:
+          - mdformat-gfm
+          - mdformat_frontmatter
+          - mdformat_footnote
+        exclude: |
+          (?x)^(
+              CHANGELOG.md|
+              README.md
+          )$
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d457b7de9..ad8f3484d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `is_torch_instance` to check against the original class of compiled models ([#8461](https://github.com/pyg-team/pytorch_geometric/pull/8461))
 - Added dense computation for `AddRandomWalkPE` ([#8431](https://github.com/pyg-team/pytorch_geometric/pull/8431))
 - Added a tutorial for point cloud processing ([#8015](https://github.com/pyg-team/pytorch_geometric/pull/8015))
 - Added `fsspec` as file system backend ([#8379](https://github.com/pyg-team/pytorch_geometric/pull/8379), [#8426](https://github.com/pyg-team/pytorch_geometric/pull/8426), [#8434](https://github.com/pyg-team/pytorch_geometric/pull/8434))
diff --git a/benchmark/citation/README.md b/benchmark/citation/README.md
index 660c89253..eae9e7de2 100644
--- a/benchmark/citation/README.md
+++ b/benchmark/citation/README.md
@@ -3,12 +3,12 @@
 Evaluation scripts for various methods on the Cora, CiteSeer and PubMed citation networks.
 Each experiment is repeated 100 times on either a fixed train/val/test split or on multiple random splits:
 
-* **[GCN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/gcn.py)**: `python gcn.py`
-* **[GAT](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/gat.py)**: `python gat.py`
-* **[Cheby](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/cheb.py)**: `python cheb.py`
-* **[SGC](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/sgc.py)**: `python sgc.py`
-* **[ARMA](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/arma.py)**: `python arma.py`
-* **[APPNP](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/appnp.py)**: `python appnp.py`
+- **[GCN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/gcn.py)**: `python gcn.py`
+- **[GAT](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/gat.py)**: `python gat.py`
+- **[Cheby](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/cheb.py)**: `python cheb.py`
+- **[SGC](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/sgc.py)**: `python sgc.py`
+- **[ARMA](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/arma.py)**: `python arma.py`
+- **[APPNP](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/citation/appnp.py)**: `python appnp.py`
 
 Run the whole test suite via
 
diff --git a/benchmark/inference/README.md b/benchmark/inference/README.md
index 6312d108a..ca0baf9e1 100644
--- a/benchmark/inference/README.md
+++ b/benchmark/inference/README.md
@@ -3,11 +3,11 @@
 ## Environment setup
 
 1. Confirm that PyG is properly installed.
-2. Install dataset package:
+1. Install dataset package:
    ```bash
    pip install ogb
    ```
-3. Install `jemalloc` for performance benchmark:
+1. Install `jemalloc` for performance benchmark:
    ```bash
    cd ${workspace}
    git clone https://github.com/jemalloc/jemalloc.git
@@ -32,11 +32,11 @@
    export LD_PRELOAD="$jemalloc_lib"
    export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:9000000000,muzzy_decay_ms:9000000000"
    ```
-2. Core binding, *e.g.*, single socket / single core / 4 cores per instance:
+1. Core binding, *e.g.*, single socket / single core / 4 cores per instance:
    ```bash
    OMP_NUM_THREADS=${CORES} numactl -C 0-${LAST_CORE} -m 0 CMD......
    ```
-3. Execute benchmarks, *e.g.*:
+1. Execute benchmarks, *e.g.*:
    ```bash
    python -u inference_benchmark.py --datasets=Reddit --models=gcn --eval-batch-sizes=512 --num-layers=2 --num-hidden-channels=64
    python -u inference_benchmark.py --datasets=Reddit --models=gcn --eval-batch-sizes=512 --num-layers=2 --num-hidden-channels=64 --use-sparse-tensor
diff --git a/benchmark/kernel/README.md b/benchmark/kernel/README.md
index 65659dcf6..a673578d7 100644
--- a/benchmark/kernel/README.md
+++ b/benchmark/kernel/README.md
@@ -3,18 +3,18 @@
 Evaluation script for various methods on [common benchmark datasets](http://graphkernels.cs.tu-dortmund.de) via 10-fold cross validation, where a training fold is randomly sampled to serve as a validation set.
 Hyperparameter selection is performed for the number of hidden units and the number of layers with respect to the validation set:
 
-* **[GCN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/gcn.py)**
-* **[GraphSAGE](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/graph_sage.py)**
-* **[GIN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/gin.py)**
-* **[Graclus](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/graclus.py)**
-* **[Top-K Pooling](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/top_k.py)**
-* **[SAG Pooling](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/sag_pool.py)**
-* **[DiffPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/diff_pool.py)**
-* **[EdgePool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/edge_pool.py)**
-* **[GlobalAttention](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/global_attention.py)**
-* **[Set2Set](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/set2set.py)**
-* **[SortPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/sort_pool.py)**
-* **[ASAPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/asap.py)**
+- **[GCN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/gcn.py)**
+- **[GraphSAGE](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/graph_sage.py)**
+- **[GIN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/gin.py)**
+- **[Graclus](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/graclus.py)**
+- **[Top-K Pooling](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/top_k.py)**
+- **[SAG Pooling](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/sag_pool.py)**
+- **[DiffPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/diff_pool.py)**
+- **[EdgePool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/edge_pool.py)**
+- **[GlobalAttention](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/global_attention.py)**
+- **[Set2Set](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/set2set.py)**
+- **[SortPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/sort_pool.py)**
+- **[ASAPool](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/kernel/asap.py)**
 
 Run (or modify) the whole test suite via
 
diff --git a/benchmark/multi_gpu/training/README.md b/benchmark/multi_gpu/training/README.md
index 56d0c5898..2b440fb44 100644
--- a/benchmark/multi_gpu/training/README.md
+++ b/benchmark/multi_gpu/training/README.md
@@ -3,6 +3,7 @@
 ## Running benchmark on CUDA GPU
 
 Run benchmark, e.g. assuming you have `n` NVIDIA GPUs:
+
 ```
 python training_benchmark_cuda.py --dataset ogbn-products --model edge_cnn --num-epochs 3 --n_gpus <n>
 ```
@@ -10,12 +11,14 @@ python training_benchmark_cuda.py --dataset ogbn-products --model edge_cnn --num
 ## Running benchmark on Intel GPU
 
 ## Environment setup
+
 ```
 install intel_extension_for_pytorch
 install oneccl_bindings_for_pytorch
 ```
 
 Run benchmark, e.g. assuming you have `n` XPUs:
+
 ```
 mpirun -np <n> python training_benchmark_xpu.py --dataset ogbn-products --model edge_cnn --num-epochs 3
 ```
diff --git a/benchmark/points/README.md b/benchmark/points/README.md
index 2e135ddf7..b9e5d65cc 100644
--- a/benchmark/points/README.md
+++ b/benchmark/points/README.md
@@ -2,8 +2,8 @@
 
 Evaluation scripts for various methods on the ModelNet10 dataset:
 
-* **[MPNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/mpnn.py)**: `python mpnn.py`
-* **[PointNet++](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/point_net.py)**: `python point_net.py`
-* **[EdgeCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/edge_cnn.py)**: `python edge_cnn.py`
-* **[SplineCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/spline_cnn.py)**: `python spline_cnn.py`
-* **[PointCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/point_cnn.py)**: `python point_cnn.py`
+- **[MPNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/mpnn.py)**: `python mpnn.py`
+- **[PointNet++](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/point_net.py)**: `python point_net.py`
+- **[EdgeCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/edge_cnn.py)**: `python edge_cnn.py`
+- **[SplineCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/spline_cnn.py)**: `python spline_cnn.py`
+- **[PointCNN](https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/points/point_cnn.py)**: `python point_cnn.py`
diff --git a/benchmark/training/README.md b/benchmark/training/README.md
index 64cc17aad..85bd6d344 100644
--- a/benchmark/training/README.md
+++ b/benchmark/training/README.md
@@ -3,11 +3,11 @@
 ## Environment setup
 
 1. Confirm that PyG is properly installed.
-2. Install dataset package:
+1. Install dataset package:
    ```bash
    pip install ogb
    ```
-3. Install `jemalloc` for performance benchmark:
+1. Install `jemalloc` for performance benchmark:
    ```bash
    cd ${workspace}
    git clone https://github.com/jemalloc/jemalloc.git
@@ -32,11 +32,11 @@
    export LD_PRELOAD="$jemalloc_lib"
    export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:9000000000,muzzy_decay_ms:9000000000"
    ```
-2. Core binding, *e.g.*, single socket / single core / 4 cores per instance:
+1. Core binding, *e.g.*, single socket / single core / 4 cores per instance:
    ```bash
    OMP_NUM_THREADS=${CORES} numactl -C 0-${LAST_CORE} -m 0 CMD......
    ```
-3. Execute benchmarks, *e.g.*:
+1. Execute benchmarks, *e.g.*:
    ```bash
    python training_benchmark.py --models=gcn --datasets=Reddit --num-workers=0 --batch-sizes=512 --num-layers=2 --num-hidden-channels=64 --num-steps=50
    python training_benchmark.py --models=gcn --datasets=Reddit --num-workers=0 --batch-sizes=512 --num-layers=2 --num-hidden-channels=64 --num-steps=50 --use-sparse-tensor
diff --git a/docker/README.md b/docker/README.md
index b91d3903d..8cedd985b 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -4,12 +4,11 @@ You can run PyG with CUDA 10.1 inside a docker image.
 The creation of [our dockerfile](https://github.com/pyg-team/pytorch_geometric/blob/master/docker/Dockerfile) refers to the dockerfiles provided by [NVIDIA](https://gitlab.com/nvidia/cuda/tree/ubuntu18.04) and [PyTorch](https://github.com/anibali/docker-pytorch).
 
 1. Download the dockerfile to your host server.
-2. `$ docker build -t "custom image name"`
-3. `$ docker run --rm -it --init --runtime=nvidia --ipc=host --network=host --volume=$PWD:/app -e NVIDIA_VISIBLE_DEVICES=0 "custom image name" /bin/bash`
+1. `$ docker build -t "custom image name"`
+1. `$ docker run --rm -it --init --runtime=nvidia --ipc=host --network=host --volume=$PWD:/app -e NVIDIA_VISIBLE_DEVICES=0 "custom image name" /bin/bash`
 
 If you encounter any problems, please feel free to contact <liangshengwen@ict.ac.cn>.
 
-
 # Singularity
 
 You can run PyG inside a singularity image. An example singularity file can be found in this folder.
@@ -49,4 +48,5 @@ to run on the CPU, or
 ```
 singularity exec --nv geometric.sif python3 gat.py
 ```
+
 to run on the GPU.
diff --git a/docs/README.md b/docs/README.md
index 3900b4980..c2a6288f5 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -3,11 +3,11 @@
 To build the documentation:
 
 1. [Build and install](https://github.com/pyg-team/pytorch_geometric/blob/master/.github/CONTRIBUTING.md#developing-pytorch-geometric) PyG from source.
-2. Install [Sphinx](https://www.sphinx-doc.org/en/master/) theme via
+1. Install [Sphinx](https://www.sphinx-doc.org/en/master/) theme via
    ```
    pip install git+https://github.com/pyg-team/pyg_sphinx_theme.git
    ```
-3. Generate the documentation file via:
+1. Generate the documentation file via:
    ```
    cd docs
    make html
diff --git a/examples/cpp/README.md b/examples/cpp/README.md
index 76a9c94bd..a4df8c88b 100644
--- a/examples/cpp/README.md
+++ b/examples/cpp/README.md
@@ -19,10 +19,10 @@ make
 Once both dependencies are sorted, we can start the CMake fun:
 
 1. Run `save_model.py` to create and save a PyG GNN model.
-2. Create a `build` directory inside the current one.
-3. From within the `build` directory, run the following commands:
-   * `cmake -DCMAKE_PREFIX_PATH="<PATH_TO_LIBTORCH>;<PATH_TO_TORCHSCATTER>;<PATH_TO_TORCH_SPARSE>" ..`
-   * `cmake --build .`
+1. Create a `build` directory inside the current one.
+1. From within the `build` directory, run the following commands:
+   - `cmake -DCMAKE_PREFIX_PATH="<PATH_TO_LIBTORCH>;<PATH_TO_TORCHSCATTER>;<PATH_TO_TORCH_SPARSE>" ..`
+   - `cmake --build .`
 
 That's it!
 You should now have a `hello-world` executable in your `build` folder.
diff --git a/examples/distributed/graphlearn_for_pytorch/README.md b/examples/distributed/graphlearn_for_pytorch/README.md
index 1f2a95ddc..85dd2d6d7 100644
--- a/examples/distributed/graphlearn_for_pytorch/README.md
+++ b/examples/distributed/graphlearn_for_pytorch/README.md
@@ -48,7 +48,7 @@ CUDA_VISIBLE_DEVICES=2,3 python dist_train_sage_supervised.py \
 **Notes:**
 
 1. You should change the `master_addr` to the IP of `node#0`.
-2. Since there is randomness during data partitioning, please ensure all nodes are using the same partitioned data when running `dist_train_sage_supervised.py`.
+1. Since there is randomness during data partitioning, please ensure all nodes are using the same partitioned data when running `dist_train_sage_supervised.py`.
 
 ### Using the `launch.py` Script
 
@@ -65,9 +65,10 @@ In distributed training (under the worker mode), each node in the cluster holds
 Thus, before the training starts, we partition the `ogbn-products` dataset into multiple partitions, each of which corresponds to a specific training worker.
 
 The partitioning occurs in three steps:
-  1. Run the partition algorithm to assign nodes to partitions.
-  2. Construct the partitioned graph structure based on the node assignment.
-  3. Split the node features and edge features into partitions.
+
+1. Run the partition algorithm to assign nodes to partitions.
+1. Construct the partitioned graph structure based on the node assignment.
+1. Split the node features and edge features into partitions.
 
 GLT supports caching graph topology and frequently accessed features in GPU to accelerate GPU sampling and feature collection.
 For feature caching, we adopt a pre-sampling-based approach to determine the hotness of nodes, and cache features for nodes with higher hotness while loading the graph.
diff --git a/examples/kuzu/papers_100M/README.md b/examples/kuzu/papers_100M/README.md
index c23bc2a97..7e30a81a7 100644
--- a/examples/kuzu/papers_100M/README.md
+++ b/examples/kuzu/papers_100M/README.md
@@ -7,7 +7,7 @@ The dataset contains approximately 100 million nodes and 1.6 billion edges.
 ## Prepare the data
 
 1. Download the dataset from [`http://snap.stanford.edu/ogb/data/nodeproppred/papers100M-bin.zip`](http://snap.stanford.edu/ogb/data/nodeproppred/papers100M-bin.zip) and put the `*.zip` file into this directory.
-2. Run `python prepare_data.py`.
+1. Run `python prepare_data.py`.
    The script will automatically extract the data and convert it to the format that Kùzu can read.
    A Kùzu database instance is then created under `papers_100M` and the data is loaded into the it.
 
diff --git a/pyproject.toml b/pyproject.toml
index b252b5168..16455501b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -151,6 +151,9 @@ filterwarnings = [
     "ignore:Sparse CSR tensor support is in beta state:UserWarning",
     "ignore:Sparse CSC tensor support is in beta state:UserWarning",
     "ignore:torch.distributed._sharded_tensor will be deprecated:DeprecationWarning",
+    # Filter `torch.compile` warnings:
+    "ignore:pkg_resources is deprecated as an API",
+    "ignore:Deprecated call to `pkg_resources.declare_namespace",
     # Filter `captum` warnings:
     "ignore:Setting backward hooks on ReLU activations:UserWarning",
     "ignore:.*did not already require gradients, required_grads has been set automatically:UserWarning",
diff --git a/torch_geometric/__init__.py b/torch_geometric/__init__.py
index 7a9ed82ae..cc3243c16 100644
--- a/torch_geometric/__init__.py
+++ b/torch_geometric/__init__.py
@@ -11,6 +11,7 @@ import torch_geometric.profile
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
 from .compile import compile
+from .isinstance import is_torch_instance
 from .debug import is_debug_enabled, debug, set_debug
 from .experimental import (is_experimental_mode_enabled, experimental_mode,
                            set_experimental_mode)
@@ -26,6 +27,7 @@ __all__ = [
     'get_home_dir',
     'set_home_dir',
     'compile',
+    'is_torch_instance',
     'is_debug_enabled',
     'debug',
     'set_debug',
diff --git a/torch_geometric/isinstance.py b/torch_geometric/isinstance.py
new file mode 100644
index 000000000..a64d74cd4
--- /dev/null
+++ b/torch_geometric/isinstance.py
@@ -0,0 +1,25 @@
+from typing import Any, Tuple, Type, Union
+
+import torch
+
+import torch_geometric.typing
+
+if torch_geometric.typing.WITH_PT20:
+    import torch._dynamo
+
+
+def is_torch_instance(obj: Any, cls: Union[Type, Tuple[Type]]) -> bool:
+    r"""Checks if the :obj:`obj` is an instance of a :obj:`cls`.
+
+    This function extends :meth:`isinstance` to be applicable during
+    :meth:`torch.compile` usage by checking against the original class of
+    compiled models.
+    """
+    # `torch.compile` removes the model inheritance and converts the model to
+    # a `torch._dynamo.OptimizedModule` instance, leading to `isinstance` being
+    # unable to check the model's inheritance. This function unwraps the
+    # compiled model before evaluating via `isinstance`.
+    if (torch_geometric.typing.WITH_PT20
+            and isinstance(obj, torch._dynamo.OptimizedModule)):
+        return isinstance(obj._orig_mod, cls)
+    return isinstance(obj, cls)


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 518269851..6e122e50c 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -1,10 +1,20 @@
 import os.path as osp
+from typing import Optional
 
 import pytest
 import torch
+from torch import Tensor
 
+import torch_geometric
 from torch_geometric.data.edge_index import EdgeIndex
-from torch_geometric.testing import onlyCUDA, withCUDA
+from torch_geometric.testing import (
+    disableExtensions,
+    onlyCUDA,
+    onlyLinux,
+    withCUDA,
+    withPackage,
+)
+from torch_geometric.utils import scatter
 
 
 def test_basic():
@@ -239,3 +249,71 @@ def test_data_loader(num_workers):
             assert isinstance(adj, EdgeIndex)
             assert adj.is_shared() == (num_workers > 0)
             assert adj._rowptr.is_shared() == (num_workers > 0)
+
+
+def test_torch_script():
+    class Model(torch.nn.Module):
+        def forward(self, x: Tensor, edge_index: EdgeIndex) -> Tensor:
+            row, col = edge_index[0], edge_index[1]
+            x_j = x[row]
+            out = scatter(x_j, col, dim_size=edge_index.num_cols)
+            return out
+
+    x = torch.randn(3, 8)
+    # Test that `num_cols` gets picked up by making last node isolated.
+    edge_index = EdgeIndex([[0, 1, 1, 2], [1, 0, 0, 1]], sparse_size=(3, 3))
+
+    model = Model()
+    expected = model(x, edge_index)
+    assert expected.size() == (3, 8)
+
+    # `torch.jit.script` does not support inheritance at the `Tensor` level :(
+    with pytest.raises(RuntimeError, match="attribute or method 'num_cols'"):
+        torch.jit.script(model)
+
+    # A valid workaround is to treat `EdgeIndex` as a regular PyTorch tensor
+    # whenever we are in script mode:
+    class ScriptableModel(torch.nn.Module):
+        def forward(self, x: Tensor, edge_index: EdgeIndex) -> Tensor:
+            row, col = edge_index[0], edge_index[1]
+            x_j = x[row]
+            dim_size: Optional[int] = None
+            if (not torch.jit.is_scripting()
+                    and isinstance(edge_index, EdgeIndex)):
+                dim_size = edge_index.num_cols
+            out = scatter(x_j, col, dim_size=dim_size)
+            return out
+
+    script_model = torch.jit.script(ScriptableModel())
+    out = script_model(x, edge_index)
+    assert out.size() == (2, 8)
+    assert torch.allclose(out, expected[:2])
+
+
+@onlyLinux
+@disableExtensions
+@withPackage('torch>=2.1.0')
+def test_compile():
+    import torch._dynamo as dynamo
+
+    class Model(torch.nn.Module):
+        def forward(self, x: Tensor, edge_index: EdgeIndex) -> Tensor:
+            row, col = edge_index[0], edge_index[1]
+            x_j = x[row]
+            out = scatter(x_j, col, dim_size=edge_index.num_cols)
+            return out
+
+    x = torch.randn(3, 8)
+    # Test that `num_cols` gets picked up by making last node isolated.
+    edge_index = EdgeIndex([[0, 1, 1, 2], [1, 0, 0, 1]], sparse_size=(3, 3))
+
+    model = Model()
+    expected = model(x, edge_index)
+    assert expected.size() == (3, 8)
+
+    explanation = dynamo.explain(model)(x, edge_index)
+    assert explanation.graph_break_count <= 0
+
+    compiled_model = torch_geometric.compile(model)
+    out = compiled_model(x, edge_index)
+    assert torch.allclose(out, expected)
diff --git a/test/test_isinstance.py b/test/test_isinstance.py
new file mode 100644
index 000000000..d2360cca7
--- /dev/null
+++ b/test/test_isinstance.py
@@ -0,0 +1,16 @@
+import torch
+
+from torch_geometric import is_torch_instance
+from torch_geometric.testing import onlyLinux, withPackage
+
+
+def test_basic():
+    assert is_torch_instance(torch.nn.Linear(1, 1), torch.nn.Linear)
+
+
+@onlyLinux
+@withPackage('torch>=2.0.0')
+def test_compile():
+    model = torch.compile(torch.nn.Linear(1, 1))
+    assert not isinstance(model, torch.nn.Linear)
+    assert is_torch_instance(model, torch.nn.Linear)