push 0.6.15 version

CMakeLists.txt

+cmake_minimum_required(VERSION 3.0)
+project(torchsparse)
+set(CMAKE_CXX_STANDARD 14)
+set(TORCHSPARSE_VERSION 0.6.15)
+
+option(WITH_CUDA "Enable CUDA support" OFF)
+option(WITH_PYTHON "Link to Python when building" ON)
+
+if(WITH_CUDA)
+  enable_language(CUDA)
+  add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
+  add_definitions(-DWITH_CUDA)
+  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
+endif()
+
+if (WITH_PYTHON)
+  add_definitions(-DWITH_PYTHON)
+  find_package(Python3 COMPONENTS Development)
+endif()
+find_package(Torch REQUIRED)
+
+file(GLOB HEADERS csrc/*.h)
+file(GLOB OPERATOR_SOURCES csrc/*.* csrc/cpu/*.*)
+if(WITH_CUDA)
+  file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} csrc/cuda/*.h csrc/cuda/*.cu)
+endif()
+
+add_library(${PROJECT_NAME} SHARED ${OPERATOR_SOURCES})
+target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES})
+if (WITH_PYTHON)
+  target_link_libraries(${PROJECT_NAME} PRIVATE Python3::Python)
+endif()
+set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchSparse)
+
+target_include_directories(${PROJECT_NAME} INTERFACE
+  $<BUILD_INTERFACE:${HEADERS}>
+  $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
+
+include(GNUInstallDirs)
+include(CMakePackageConfigHelpers)
+
+set(PHMAP_DIR third_party/parallel-hashmap)
+target_include_directories(${PROJECT_NAME} PRIVATE ${PHMAP_DIR})
+
+set(TORCHSPARSE_CMAKECONFIG_INSTALL_DIR "share/cmake/TorchSparse" CACHE STRING "install path for TorchSparseConfig.cmake")
+
+configure_package_config_file(cmake/TorchSparseConfig.cmake.in
+  "${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfig.cmake"
+  INSTALL_DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
+
+write_basic_package_version_file(${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfigVersion.cmake
+	VERSION ${TORCHSPARSE_VERSION}
+  COMPATIBILITY AnyNewerVersion)
+
+install(FILES ${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfig.cmake
+  ${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfigVersion.cmake
+  DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
+
+install(TARGETS ${PROJECT_NAME}
+  EXPORT TorchSparseTargets
+  LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  )
+
+install(EXPORT TorchSparseTargets
+  NAMESPACE TorchSparse::
+  DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
+
+install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
+install(FILES
+  csrc/cpu/convert_cpu.h
+  csrc/cpu/diag_cpu.h
+  csrc/cpu/metis_cpu.h
+  csrc/cpu/rw_cpu.h
+  csrc/cpu/saint_cpu.h
+  csrc/cpu/sample_cpu.h
+  csrc/cpu/spmm_cpu.h
+  csrc/cpu/spspmm_cpu.h
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
+if(WITH_CUDA)
+  install(FILES
+    csrc/cuda/convert_cuda.h
+    csrc/cuda/diag_cuda.h
+    csrc/cuda/rw_cuda.h
+    csrc/cuda/spmm_cuda.h
+    csrc/cuda/spspmm_cuda.h
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
+endif()
+
+if(WITH_CUDA)
+  set_property(TARGET torch_cuda PROPERTY INTERFACE_COMPILE_OPTIONS "")
+  set_property(TARGET torch_cpu PROPERTY INTERFACE_COMPILE_OPTIONS "")
+endif()

README.md

-# <div align="center"><strong>torch-sparse-0.6.13</strong></div>
-## 简介
-torch-sparse是PyTorch的一个扩展库，用于处理稀疏数据。它提供了一些功能强大的稀疏矩阵操作，以及能够高效执行稀疏计算的函数和工具。
+[pypi-image]: https://badge.fury.io/py/torch-sparse.svg
+[pypi-url]: https://pypi.python.org/pypi/torch-sparse
+[testing-image]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/testing.yml/badge.svg
+[testing-url]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/testing.yml
+[linting-image]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/linting.yml/badge.svg
+[linting-url]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/linting.yml
+[coverage-image]: https://codecov.io/gh/rusty1s/pytorch_sparse/branch/master/graph/badge.svg
+[coverage-url]: https://codecov.io/github/rusty1s/pytorch_sparse?branch=master
 
-## 依赖安装
-+ pytorch1.10或者pytorch1.13 以及对应的torchvision（建议dtk-22.04.2、dtk-23.04与dtk-23.10）
-+ python 3.7-3.10
+# PyTorch Sparse
 
-### 1、使用源码编译方式安装
+[![PyPI Version][pypi-image]][pypi-url]
+[![Testing Status][testing-image]][testing-url]
+[![Linting Status][linting-image]][linting-url]
+[![Code Coverage][coverage-image]][coverage-url]
 
-#### 编译环境准备
-提供2种环境准备方式：
+--------------------------------------------------------------------------------
 
-1. 基于光源pytorch基础镜像环境：镜像下载地址：[https://sourcefind.cn/#/image/dcu/pytorch](https://sourcefind.cn/#/image/dcu/pytorch)，根据pytorch、python、dtk及系统下载对应的镜像版本。
+This package consists of a small extension library of optimized sparse matrix operations with autograd support.
+This package currently consists of the following methods:
 
-2. 基于现有python环境：安装pytorch和torchvision，whl包下载目录：[https://cancon.hpccube.com:65024/4/main/pytorch](https://cancon.hpccube.com:65024/4/main/pytorch)，[https://cancon.hpccube.com:65024/4/main/vision](https://cancon.hpccube.com:65024/4/main/vision)，根据python、dtk版本,下载对应pytorch和torchvision的whl包。安装命令如下：
-```shell
-pip install torch* (下载的torch的whl包)
-pip install torchvision* (下载的torchvision的whl包)
-pip install setuptools==59.5.0 wheel
+* **[Coalesce](#coalesce)**
+* **[Transpose](#transpose)**
+* **[Sparse Dense Matrix Multiplication](#sparse-dense-matrix-multiplication)**
+* **[Sparse Sparse Matrix Multiplication](#sparse-sparse-matrix-multiplication)**
+
+All included operations work on varying data types and are implemented both for CPU and GPU.
+To avoid the hazzle of creating [`torch.sparse_coo_tensor`](https://pytorch.org/docs/stable/torch.html?highlight=sparse_coo_tensor#torch.sparse_coo_tensor), this package defines operations on sparse tensors by simply passing `index` and `value` tensors as arguments ([with same shapes as defined in PyTorch](https://pytorch.org/docs/stable/sparse.html)).
+Note that only `value` comes with autograd support, as `index` is discrete and therefore not differentiable.
+
+## Installation
+
+### Anaconda
+
+**Update:** You can now install `pytorch-sparse` via [Anaconda](https://anaconda.org/pyg/pytorch-sparse) for all major OS/PyTorch/CUDA combinations 🤗
+Given that you have [`pytorch >= 1.8.0` installed](https://pytorch.org/get-started/locally/), simply run
+
+```
+conda install pytorch-sparse -c pyg
 ```
+#### Note: Conda packages are not published for PyTorch 1.12 yet
+### Binaries
+
+We alternatively provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://data.pyg.org/whl).
+
+#### PyTorch 1.12
+
+To install the binaries for PyTorch 1.12.0, simply run
 
-#### 源码编译安装
-- 代码下载
-```shell
-git clone http://developer.hpccube.com/codes/aicomponent/torch-sparce # 根据编译需要切换分支
 ```
-- 源码编译（进入torch-sparse目录）：
+pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.12.0+${CUDA}.html
 ```
-export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$C_INCLUDE_PATH
-export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$CPLUS_INCLUDE_PATH
-export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$C_INCLUDE_PATH
-export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$CPLUS_INCLUDE_PATH
-export C_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$C_INCLUDE_PATH
-export CPLUS_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$CPLUS_INCLUDE_PATH
-export LD_LIBRARY_PATH=$ROCM_PATH/rocrand/lib:$LD_LIBRARY_PATH
-export FORCE_ONLY_HIP=1
-export CC=hipcc
-export CXX=hipcc
 
-python setup.py install
+where `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, or `cu116` depending on your PyTorch installation.
+
+|             | `cpu` | `cu102` | `cu113` | `cu116` |
+|-------------|-------|---------|---------|---------|
+| **Linux**   | ✅    | ✅      | ✅      | ✅      |
+| **Windows** | ✅    |         | ✅      | ✅      |
+| **macOS**   | ✅    |         |         |         |
+
+#### PyTorch 1.11
+
+To install the binaries for PyTorch 1.11.0, simply run
+
+```
+pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.11.0+${CUDA}.html
+```
+
+where `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, or `cu115` depending on your PyTorch installation.
+
+|             | `cpu` | `cu102` | `cu113` | `cu115` |
+|-------------|-------|---------|---------|---------|
+| **Linux**   | ✅    | ✅      | ✅      | ✅      |
+| **Windows** | ✅    |         | ✅      | ✅      |
+| **macOS**   | ✅    |         |         |         |
+
+**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, and PyTorch 1.10.0/1.10.1/1.10.2 (following the same procedure).
+For older versions, you might need to explicitly specify the latest supported version number in order to prevent a manual installation from source.
+You can look up the latest supported version number [here](https://data.pyg.org/whl).
+
+### From source
+
+Ensure that at least PyTorch 1.7.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*:
+
+```
+$ python -c "import torch; print(torch.__version__)"
+>>> 1.7.0
+
+$ echo $PATH
+>>> /usr/local/cuda/bin:...
+
+$ echo $CPATH
+>>> /usr/local/cuda/include:...
+```
+
+If you want to additionally build `torch-sparse` with METIS support, *e.g.* for partioning, please download and install the [METIS library](https://web.archive.org/web/20211119110155/http://glaros.dtc.umn.edu/gkhome/metis/metis/download) by following the instructions in the `Install.txt` file.
+Note that METIS needs to be installed with 64 bit `IDXTYPEWIDTH` by changing `include/metis.h`.
+Afterwards, set the environment variable `WITH_METIS=1`.
+
+Then run:
+
+```
+pip install torch-scatter torch-sparse
 ```
-#### 注意事项
-+ 若使用pip install下载安装过慢，可添加pypi清华源：-i https://pypi.tuna.tsinghua.edu.cn/simple/
-+ ROCM_PATH为dtk的路径，默认为/opt/dtk
 
-## 验证
+When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
+In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:
+
+```
+export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
+```
+
+## Functions
+
+### Coalesce
+
+```
+torch_sparse.coalesce(index, value, m, n, op="add") -> (torch.LongTensor, torch.Tensor)
+```
+
+Row-wise sorts `index` and removes duplicate entries.
+Duplicate entries are removed by scattering them together.
+For scattering, any operation of [`torch_scatter`](https://github.com/rusty1s/pytorch_scatter) can be used.
+
+#### Parameters
+
+* **index** *(LongTensor)* - The index tensor of sparse matrix.
+* **value** *(Tensor)* - The value tensor of sparse matrix.
+* **m** *(int)* - The first dimension of sparse matrix.
+* **n** *(int)* - The second dimension of sparse matrix.
+* **op** *(string, optional)* - The scatter operation to use. (default: `"add"`)
+
+#### Returns
+
+* **index** *(LongTensor)* - The coalesced index tensor of sparse matrix.
+* **value** *(Tensor)* - The coalesced value tensor of sparse matrix.
+
+#### Example
 
 ```python
 import torch
@@ -68,10 +160,156 @@ tensor([[6.0, 8.0],
         [5.0, 6.0]])
 ```
 
-## Known Issue
-- 该库没有基于cpu环境修改，仅支持dcu，请在有dcu卡的环境运行。
-- 如需完整使用所有pyg功能，请pip install torch-geometric
+### Transpose
 
-## 参考资料
-- [README_ORIGIN](README_ORIGIN.md)
-- [https://pypi.org/project/torch-sparse/0.6.13/](https://pypi.org/project/torch-sparse/0.6.13/)
+```
+torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
+```
+
+Transposes dimensions 0 and 1 of a sparse matrix.
+
+#### Parameters
+
+* **index** *(LongTensor)* - The index tensor of sparse matrix.
+* **value** *(Tensor)* - The value tensor of sparse matrix.
+* **m** *(int)* - The first dimension of sparse matrix.
+* **n** *(int)* - The second dimension of sparse matrix.
+* **coalesced** *(bool, optional)* - If set to `False`, will not coalesce the output. (default: `True`)
+
+#### Returns
+
+* **index** *(LongTensor)* - The transposed index tensor of sparse matrix.
+* **value** *(Tensor)* - The transposed value tensor of sparse matrix.
+
+#### Example
+
+```python
+import torch
+from torch_sparse import transpose
+
+index = torch.tensor([[1, 0, 1, 0, 2, 1],
+                      [0, 1, 1, 1, 0, 0]])
+value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
+
+index, value = transpose(index, value, 3, 2)
+```
+
+```
+print(index)
+tensor([[0, 0, 1, 1],
+        [1, 2, 0, 1]])
+print(value)
+tensor([[7.0, 9.0],
+        [5.0, 6.0],
+        [6.0, 8.0],
+        [3.0, 4.0]])
+```
+
+### Sparse Dense Matrix Multiplication
+
+```
+torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
+```
+
+Matrix product of a sparse matrix with a dense matrix.
+
+#### Parameters
+
+* **index** *(LongTensor)* - The index tensor of sparse matrix.
+* **value** *(Tensor)* - The value tensor of sparse matrix.
+* **m** *(int)* - The first dimension of sparse matrix.
+* **n** *(int)* - The second dimension of sparse matrix.
+* **matrix** *(Tensor)* - The dense matrix.
+
+#### Returns
+
+* **out** *(Tensor)* - The dense output matrix.
+
+#### Example
+
+```python
+import torch
+from torch_sparse import spmm
+
+index = torch.tensor([[0, 0, 1, 2, 2],
+                      [0, 2, 1, 0, 1]])
+value = torch.Tensor([1, 2, 4, 1, 3])
+matrix = torch.Tensor([[1, 4], [2, 5], [3, 6]])
+
+out = spmm(index, value, 3, 3, matrix)
+```
+
+```
+print(out)
+tensor([[7.0, 16.0],
+        [8.0, 20.0],
+        [7.0, 19.0]])
+```
+
+### Sparse Sparse Matrix Multiplication
+
+```
+torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n) -> (torch.LongTensor, torch.Tensor)
+```
+
+Matrix product of two sparse tensors.
+Both input sparse matrices need to be **coalesced** (use the `coalesced` attribute to force).
+
+#### Parameters
+
+* **indexA** *(LongTensor)* - The index tensor of first sparse matrix.
+* **valueA** *(Tensor)* - The value tensor of first sparse matrix.
+* **indexB** *(LongTensor)* - The index tensor of second sparse matrix.
+* **valueB** *(Tensor)* - The value tensor of second sparse matrix.
+* **m** *(int)* - The first dimension of first sparse matrix.
+* **k** *(int)* - The second dimension of first sparse matrix and first dimension of second sparse matrix.
+* **n** *(int)* - The second dimension of second sparse matrix.
+* **coalesced** *(bool, optional)*: If set to `True`, will coalesce both input sparse matrices. (default: `False`)
+
+#### Returns
+
+* **index** *(LongTensor)* - The output index tensor of sparse matrix.
+* **value** *(Tensor)* - The output value tensor of sparse matrix.
+
+#### Example
+
+```python
+import torch
+from torch_sparse import spspmm
+
+indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
+valueA = torch.Tensor([1, 2, 3, 4, 5])
+
+indexB = torch.tensor([[0, 2], [1, 0]])
+valueB = torch.Tensor([2, 4])
+
+indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)
+```
+
+```
+print(indexC)
+tensor([[0, 1, 2],
+        [0, 1, 1]])
+print(valueC)
+tensor([8.0, 6.0, 8.0])
+```
+
+## Running tests
+
+```
+pytest
+```
+
+## C++ API
+
+`torch-sparse` also offers a C++ API that contains C++ equivalent of python models.
+For this, we need to add `TorchLib` to the `-DCMAKE_PREFIX_PATH` (*e.g.*, it may exists in `{CONDA}/lib/python{X.X}/site-packages/torch` if installed via `conda`):
+
+```
+mkdir build
+cd build
+# Add -DWITH_CUDA=on support for CUDA support
+cmake -DCMAKE_PREFIX_PATH="..." ..
+make
+make install
+```

benchmark/.gitignore

+*.mat
+*.tmp

benchmark/main.py

+import time
+import os.path as osp
+import itertools
+
+import argparse
+import wget
+import torch
+from scipy.io import loadmat
+
+from torch_scatter import scatter_add
+from torch_sparse.tensor import SparseTensor
+
+short_rows = [
+    ('DIMACS10', 'citationCiteseer'),
+    ('SNAP', 'web-Stanford'),
+]
+long_rows = [
+    ('Janna', 'StocF-1465'),
+    ('GHS_psdef', 'ldoor'),
+]
+
+
+def download(dataset):
+    url = 'https://sparse.tamu.edu/mat/{}/{}.mat'
+    for group, name in itertools.chain(long_rows, short_rows):
+        if not osp.exists(f'{name}.mat'):
+            print(f'Downloading {group}/{name}:')
+            wget.download(url.format(group, name))
+            print('')
+
+
+def bold(text, flag=True):
+    return f'\033[1m{text}\033[0m' if flag else text
+
+
+@torch.no_grad()
+def correctness(dataset):
+    group, name = dataset
+    mat_scipy = loadmat(f'{name}.mat')['Problem'][0][0][2].tocsr()
+    row = torch.from_numpy(mat_scipy.tocoo().row).to(args.device, torch.long)
+    col = torch.from_numpy(mat_scipy.tocoo().col).to(args.device, torch.long)
+    mat = SparseTensor(row=row, col=col, sparse_sizes=mat_scipy.shape)
+    mat.fill_cache_()
+    mat_pytorch = mat.to_torch_sparse_coo_tensor().coalesce()
+
+    for size in sizes:
+        try:
+            x = torch.randn((mat.size(1), size), device=args.device)
+
+            out1 = mat @ x
+            out2 = mat_pytorch @ x
+
+            assert torch.allclose(out1, out2, atol=1e-4)
+
+        except RuntimeError as e:
+            if 'out of memory' not in str(e):
+                raise RuntimeError(e)
+            torch.cuda.empty_cache()
+
+
+def time_func(func, x):
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t = time.perf_counter()
+
+        if not args.with_backward:
+            with torch.no_grad():
+                for _ in range(iters):
+                    func(x)
+        else:
+            x = x.requires_grad_()
+            for _ in range(iters):
+                out = func(x)
+                out = out[0] if isinstance(out, tuple) else out
+                torch.autograd.grad(out, x, out, only_inputs=True)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        return time.perf_counter() - t
+    except RuntimeError as e:
+        if 'out of memory' not in str(e):
+            raise RuntimeError(e)
+        torch.cuda.empty_cache()
+        return float('inf')
+
+
+def timing(dataset):
+    group, name = dataset
+    mat_scipy = loadmat(f'{name}.mat')['Problem'][0][0][2].tocsr()
+    row = torch.from_numpy(mat_scipy.tocoo().row).to(args.device, torch.long)
+    col = torch.from_numpy(mat_scipy.tocoo().col).to(args.device, torch.long)
+    mat = SparseTensor(row=row, col=col, sparse_sizes=mat_scipy.shape)
+    mat.fill_cache_()
+    mat_pytorch = mat.to_torch_sparse_coo_tensor().coalesce()
+    mat_scipy = mat.to_scipy(layout='csr')
+
+    def scatter(x):
+        return scatter_add(x[col], row, dim=0, dim_size=mat_scipy.shape[0])
+
+    def spmm_scipy(x):
+        if x.is_cuda:
+            raise RuntimeError('out of memory')
+        return mat_scipy @ x
+
+    def spmm_pytorch(x):
+        return mat_pytorch @ x
+
+    def spmm(x):
+        return mat @ x
+
+    t1, t2, t3, t4 = [], [], [], []
+
+    for size in sizes:
+        try:
+            x = torch.randn((mat.size(1), size), device=args.device)
+
+            t1 += [time_func(scatter, x)]
+            t2 += [time_func(spmm_scipy, x)]
+            t3 += [time_func(spmm_pytorch, x)]
+            t4 += [time_func(spmm, x)]
+
+            del x
+
+        except RuntimeError as e:
+            if 'out of memory' not in str(e):
+                raise RuntimeError(e)
+            torch.cuda.empty_cache()
+            for t in (t1, t2, t3, t4):
+                t.append(float('inf'))
+
+    ts = torch.tensor([t1, t2, t3, t4])
+    winner = torch.zeros_like(ts, dtype=torch.bool)
+    winner[ts.argmin(dim=0), torch.arange(len(sizes))] = 1
+    winner = winner.tolist()
+
+    name = f'{group}/{name}'
+    print(f'{bold(name)} (avg row length: {mat.avg_row_length():.2f}):')
+    print('\t'.join(['            '] + [f'{size:>5}' for size in sizes]))
+    print('\t'.join([bold('Scatter     ')] +
+                    [bold(f'{t:.5f}', f) for t, f in zip(t1, winner[0])]))
+    print('\t'.join([bold('SPMM SciPy  ')] +
+                    [bold(f'{t:.5f}', f) for t, f in zip(t2, winner[1])]))
+    print('\t'.join([bold('SPMM PyTorch')] +
+                    [bold(f'{t:.5f}', f) for t, f in zip(t3, winner[2])]))
+    print('\t'.join([bold('SPMM Own    ')] +
+                    [bold(f'{t:.5f}', f) for t, f in zip(t4, winner[3])]))
+    print()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--with_backward', action='store_true')
+    parser.add_argument('--device', type=str, default='cuda')
+    args = parser.parse_args()
+    iters = 1 if args.device == 'cpu' else 20
+    sizes = [1, 16, 32, 64, 128, 256, 512]
+    sizes = sizes[:4] if args.device == 'cpu' else sizes
+
+    for _ in range(10):  # Warmup.
+        torch.randn(100, 100, device=args.device).sum()
+    for dataset in itertools.chain(short_rows, long_rows):
+        download(dataset)
+        correctness(dataset)
+        timing(dataset)

benchmark/ptr2ind.py

+import os
+import wget
+import time
+import errno
+import argparse
+import os.path as osp
+
+import torch
+from scipy.io import loadmat
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--root', type=str, default='/tmp/test_ptr2ind')
+args = parser.parse_args()
+
+matrices = [
+    ('DIMACS10', 'citationCiteseer'),
+    ('SNAP', 'web-Stanford'),
+    ('Janna', 'StocF-1465'),
+    ('GHS_psdef', 'ldoor'),
+]
+
+
+def get_torch_sparse_coo_tensor(root, group, name):
+    path = osp.join(root, f'{name}.mat')
+    if not osp.exists(path):
+        try:
+            os.makedirs(root)
+        except OSError as e:
+            if e.errno != errno.EEXIST and osp.isdir(path):
+                raise e
+        url = f'https://sparse.tamu.edu/mat/{group}/{name}.mat'
+        print(f'Downloading {group}/{name}:')
+        wget.download(url, path)
+    matrix = loadmat(path)['Problem'][0][0][2].tocoo()
+    row = torch.from_numpy(matrix.row).to(torch.long)
+    col = torch.from_numpy(matrix.col).to(torch.long)
+    index = torch.stack([row, col], dim=0)
+    value = torch.from_numpy(matrix.data).to(torch.float)
+    print(f'{name}.mat: shape={matrix.shape} nnz={row.numel()}')
+    return torch.sparse_coo_tensor(index, value, matrix.shape).coalesce()
+
+
+def time_func(matrix, op, duration=5.0, warmup=1.0):
+    t = time.time()
+    while (time.time() - t) < warmup:
+        op(matrix)
+    torch.cuda.synchronize()
+    count = 0
+    t = time.time()
+    while (time.time() - t) < duration:
+        op(matrix)
+        count += 1
+        torch.cuda.synchronize()
+    return (time.time() - t) / count
+
+
+def bucketize(matrix):
+    row_indices = matrix.indices()[0]
+    arange = torch.arange(matrix.size(0) + 1, device=row_indices.device)
+    return torch.bucketize(arange, row_indices)
+
+
+def convert_coo_to_csr(matrix):
+    row_indices = matrix.indices()[0]
+    return torch._convert_coo_to_csr(row_indices, matrix.size(0))
+
+
+for device in ['cpu', 'cuda']:
+    print('DEVICE:', device)
+    for group, name in matrices:
+        matrix = get_torch_sparse_coo_tensor(args.root, group, name)
+        matrix = matrix.to(device)
+
+        out1 = bucketize(matrix)
+        out2 = convert_coo_to_csr(matrix)
+        assert out1.tolist() == out2.tolist()
+
+        t = time_func(matrix, bucketize, duration=5.0, warmup=1.0)
+        print('old impl:', t)
+        t = time_func(matrix, convert_coo_to_csr, duration=5.0, warmup=1.0)
+        print('new impl:', t)
+        print()

cmake/TorchSparseConfig.cmake.in

+# TorchSparseConfig.cmake
+# --------------------
+#
+# Exported targets:: Sparse
+#
+
+@PACKAGE_INIT@
+
+set(PN TorchSparse)
+set(${PN}_INCLUDE_DIR "${PACKAGE_PREFIX_DIR}/@CMAKE_INSTALL_INCLUDEDIR@")
+set(${PN}_LIBRARY "")
+set(${PN}_DEFINITIONS USING_${PN})
+
+check_required_components(${PN})
+
+
+if(NOT (CMAKE_VERSION VERSION_LESS 3.0))
+#-----------------------------------------------------------------------------
+# Don't include targets if this file is being picked up by another
+# project which has already built this as a subproject
+#-----------------------------------------------------------------------------
+if(NOT TARGET ${PN}::TorchSparse)
+include("${CMAKE_CURRENT_LIST_DIR}/${PN}Targets.cmake")
+
+if(NOT TARGET torch_library)
+find_package(Torch REQUIRED)
+endif()
+if(NOT TARGET Python3::Python)
+find_package(Python3 COMPONENTS Development)
+endif()
+target_link_libraries(TorchSparse::TorchSparse INTERFACE ${TORCH_LIBRARIES} Python3::Python)
+
+if(@WITH_CUDA@)
+  target_compile_definitions(TorchSparse::TorchSparse INTERFACE WITH_CUDA)
+endif()
+
+endif()
+endif()

conda/pytorch-sparse/README.md

+```
+./build_conda.sh 3.9 1.12.0 cu113  # python, pytorch and cuda version
+```

conda/pytorch-sparse/bld.bat

+copy "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC\\Tools\\MSVC\\14.29.30133\\lib\\x64\\metis.lib" %LIBRARY_LIB%
+if errorlevel 1 exit 1
+copy "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC\\Tools\\MSVC\\14.29.30133\\include\\metis.h" %LIBRARY_INC%
+if errorlevel 1 exit 1
+
+"%PYTHON%" -m pip install .
+if errorlevel 1 exit 1

conda/pytorch-sparse/build.sh

+$PYTHON -m pip install .

conda/pytorch-sparse/build_conda.sh

+#!/bin/bash
+
+export PYTHON_VERSION=$1
+export TORCH_VERSION=$2
+export CUDA_VERSION=$3
+
+export CONDA_PYTORCH_CONSTRAINT="pytorch==${TORCH_VERSION%.*}.*"
+
+if [ "${CUDA_VERSION}" = "cpu" ]; then
+  export CONDA_CUDATOOLKIT_CONSTRAINT="cpuonly  # [not osx]"
+else
+  case $CUDA_VERSION in
+    cu116)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.6.*"
+      ;;
+    cu115)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.5.*"
+      ;;
+    cu113)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.3.*"
+      ;;
+    cu111)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.1.*"
+      ;;
+    cu102)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==10.2.*"
+      ;;
+    cu101)
+      export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==10.1.*"
+      ;;
+    *)
+      echo "Unrecognized CUDA_VERSION=$CUDA_VERSION"
+      exit 1
+      ;;
+  esac
+fi
+
+echo "PyTorch $TORCH_VERSION+$CUDA_VERSION"
+echo "- $CONDA_PYTORCH_CONSTRAINT"
+echo "- $CONDA_CUDATOOLKIT_CONSTRAINT"
+
+if [ "${CUDA_VERSION}" = "cu116" ]; then
+  conda build . -c pytorch -c default -c nvidia -c conda-forge --output-folder "$HOME/conda-bld"
+else
+  conda build . -c pytorch -c default -c nvidia --output-folder "$HOME/conda-bld"
+fi

conda/pytorch-sparse/meta.yaml

+package:
+  name: pytorch-sparse
+  version: 0.6.15
+
+source:
+  path: ../..
+
+requirements:
+  build:
+    - {{ compiler('c') }}  # [win]
+
+  host:
+    - pip
+    - python {{ environ.get('PYTHON_VERSION') }}
+    - {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
+    - {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
+
+  run:
+    - scipy
+    - pytorch-scatter
+    - python {{ environ.get('PYTHON_VERSION') }}
+    - {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
+    - {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
+
+build:
+  string: py{{ environ.get('PYTHON_VERSION').replace('.', '') }}_torch_{{ environ['TORCH_VERSION'] }}_{{ environ['CUDA_VERSION'] }}
+  script_env:
+    - FORCE_CUDA
+    - TORCH_CUDA_ARCH_LIST
+    - WITH_METIS=1
+  preserve_egg_dir: True
+
+test:
+  imports:
+    - torch_sparse
+
+about:
+  home: https://github.com/rusty1s/pytorch_sparse
+  license: MIT
+  summary: PyTorch Extension Library of Optimized Autograd Sparse Matrix Operations

csrc/convert.cpp

@@ -5,13 +5,13 @@
 
 #include "cpu/convert_cpu.h"
 
-#ifdef WITH_HIP
-#include "hip/convert_hip.h"
+#ifdef WITH_CUDA
+#include "cuda/convert_cuda.h"
 #endif
 
 #ifdef _WIN32
 #ifdef WITH_PYTHON
-#ifdef WITH_HIP
+#ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__convert_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
@@ -21,7 +21,7 @@ PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
 
 SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
   if (ind.device().is_cuda()) {
-#ifdef WITH_HIP
+#ifdef WITH_CUDA
     return ind2ptr_cuda(ind, M);
 #else
     AT_ERROR("Not compiled with CUDA support");
@@ -33,7 +33,7 @@ SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
 
 SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) {
   if (ptr.device().is_cuda()) {
-#ifdef WITH_HIP
+#ifdef WITH_CUDA
     return ptr2ind_cuda(ptr, E);
 #else
     AT_ERROR("Not compiled with CUDA support");

csrc/cpu/diag_cpu.cpp

@@ -13,7 +13,7 @@ torch::Tensor non_diag_mask_cpu(torch::Tensor row, torch::Tensor col, int64_t M,
   auto row_data = row.data_ptr<int64_t>();
   auto col_data = col.data_ptr<int64_t>();
 
-  auto mask = torch::zeros(E + num_diag, row.options().dtype(torch::kBool));
+  auto mask = torch::zeros({E + num_diag}, row.options().dtype(torch::kBool));
   auto mask_data = mask.data_ptr<bool>();
 
   int64_t r, c;

csrc/cpu/ego_sample_cpu.cpp

@@ -19,8 +19,6 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col,
                          torch::Tensor idx, int64_t depth,
                          int64_t num_neighbors, bool replace) {
 
-  srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
-
   std::vector<torch::Tensor> out_rowptrs(idx.numel() + 1);
   std::vector<torch::Tensor> out_cols(idx.numel());
   std::vector<torch::Tensor> out_n_ids(idx.numel());
@@ -56,14 +54,14 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col,
             }
           } else if (replace) {
             for (int64_t j = 0; j < num_neighbors; j++) {
-              w = col_data[row_start + (rand() % row_count)];
+              w = col_data[row_start + uniform_randint(row_count)];
               n_id_set.insert(w);
               n_ids.push_back(w);
             }
           } else {
             std::unordered_set<int64_t> perm;
             for (int64_t j = row_count - num_neighbors; j < row_count; j++) {
-              if (!perm.insert(rand() % j).second) {
+              if (!perm.insert(uniform_randint(j)).second) {
                 perm.insert(j);
               }
             }

csrc/cpu/hgt_sample_cpu.cpp

@@ -105,8 +105,6 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
                const c10::Dict<node_t, vector<int64_t>> &num_samples_dict,
                const int64_t num_hops) {
 
-  srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
-
   // Create a mapping to convert single string relations to edge type triplets:
   unordered_map<rel_t, edge_t> to_edge_type;
   for (const auto &kv : colptr_dict) {
@@ -226,12 +224,10 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
         }
       }
     }
-    if (rows.size() > 0) {
     out_row_dict.insert(rel_type, from_vector<int64_t>(rows));
     out_col_dict.insert(rel_type, from_vector<int64_t>(cols));
     out_edge_dict.insert(rel_type, from_vector<int64_t>(edges));
-    }
-  }
+}
 
   // Generate tensor-valued output node dictionary (line 20):
   for (const auto &kv : nodes_dict) {

csrc/cpu/metis_cpu.cpp

@@ -44,7 +44,7 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
     vwgt = optional_node_weight.value().data_ptr<int64_t>();
 
   int64_t objval = -1;
-  auto part = torch::empty(nvtxs, rowptr.options());
+  auto part = torch::empty({nvtxs}, rowptr.options());
   auto part_data = part.data_ptr<int64_t>();
 
   if (recursive) {
@@ -99,7 +99,7 @@ mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
 
   mtmetis_pid_type nparts = num_parts;
   mtmetis_wgt_type objval = -1;
-  auto part = torch::empty(nvtxs, rowptr.options());
+  auto part = torch::empty({nvtxs}, rowptr.options());
   mtmetis_pid_type *part_data = (mtmetis_pid_type *)part.data_ptr<int64_t>();
 
   double *opts = mtmetis_init_options();

csrc/cpu/neighbor_sample_cpu.cpp

@@ -10,16 +10,16 @@ using namespace std;
 
 namespace {
 
+typedef phmap::flat_hash_map<pair<int64_t, int64_t>, int64_t> temporarl_edge_dict;
+
 template <bool replace, bool directed>
 tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 sample(const torch::Tensor &colptr, const torch::Tensor &row,
        const torch::Tensor &input_node, const vector<int64_t> num_neighbors) {
 
-  srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
-
   // Initialize some data structures for the sampling process:
   vector<int64_t> samples;
-  unordered_map<int64_t, int64_t> to_local_node;
+  phmap::flat_hash_map<int64_t, int64_t> to_local_node;
 
   auto *colptr_data = colptr.data_ptr<int64_t>();
   auto *row_data = row.data_ptr<int64_t>();
@@ -59,7 +59,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
         }
       } else if (replace) {
         for (int64_t j = 0; j < num_samples; j++) {
-          const int64_t offset = col_start + rand() % col_count;
+          const int64_t offset = col_start + uniform_randint(col_count);
           const int64_t &v = row_data[offset];
           const auto res = to_local_node.insert({v, samples.size()});
           if (res.second)
@@ -73,7 +73,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
       } else {
         unordered_set<int64_t> rnd_indices;
         for (int64_t j = col_count - num_samples; j < col_count; j++) {
-          int64_t rnd = rand() % j;
+          int64_t rnd = uniform_randint(j);
           if (!rnd_indices.insert(rnd).second) {
             rnd = j;
             rnd_indices.insert(j);
@@ -95,7 +95,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
   }
 
   if (!directed) {
-    unordered_map<int64_t, int64_t>::iterator iter;
+    phmap::flat_hash_map<int64_t, int64_t>::iterator iter;
     for (int64_t i = 0; i < (int64_t)samples.size(); i++) {
       const auto &w = samples[i];
       const auto &col_start = colptr_data[w];
@@ -116,7 +116,20 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
                     from_vector<int64_t>(cols), from_vector<int64_t>(edges));
 }
 
-template <bool replace, bool directed>
+inline bool satisfy_time(const c10::Dict<node_t, torch::Tensor> &node_time_dict,
+                         const node_t &src_node_type, int64_t dst_time,
+                         int64_t src_node) {
+  try {
+    // Check whether src -> dst obeys the time constraint
+    const torch::Tensor &src_node_time = node_time_dict.at(src_node_type);
+    return src_node_time.data_ptr<int64_t>()[src_node] <= dst_time;
+  } catch (const std::out_of_range& e) {
+    // If no time is given, fall back to normal sampling
+    return true;
+  }
+}
+
+template <bool replace, bool directed, bool temporal>
 tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
       c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
 hetero_sample(const vector<node_t> &node_types,
@@ -125,24 +138,29 @@ hetero_sample(const vector<node_t> &node_types,
               const c10::Dict<rel_t, torch::Tensor> &row_dict,
               const c10::Dict<node_t, torch::Tensor> &input_node_dict,
               const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
+              const c10::Dict<node_t, torch::Tensor> &node_time_dict,
               const int64_t num_hops) {
 
-  srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
-
   // Create a mapping to convert single string relations to edge type triplets:
-  unordered_map<rel_t, edge_t> to_edge_type;
+  phmap::flat_hash_map<rel_t, edge_t> to_edge_type;
   for (const auto &k : edge_types)
     to_edge_type[get<0>(k) + "__" + get<1>(k) + "__" + get<2>(k)] = k;
 
   // Initialize some data structures for the sampling process:
-  unordered_map<node_t, vector<int64_t>> samples_dict;
-  unordered_map<node_t, unordered_map<int64_t, int64_t>> to_local_node_dict;
+  phmap::flat_hash_map<node_t, vector<int64_t>> samples_dict;
+  phmap::flat_hash_map<node_t, vector<pair<int64_t, int64_t>>> temp_samples_dict;
+  phmap::flat_hash_map<node_t, phmap::flat_hash_map<int64_t, int64_t>> to_local_node_dict;
+  phmap::flat_hash_map<node_t, temporarl_edge_dict> temp_to_local_node_dict;
+  phmap::flat_hash_map<node_t, vector<int64_t>> root_time_dict;
   for (const auto &node_type : node_types) {
     samples_dict[node_type];
+    temp_samples_dict[node_type];
     to_local_node_dict[node_type];
+    temp_to_local_node_dict[node_type];
+    root_time_dict[node_type];
   }
 
-  unordered_map<rel_t, vector<int64_t>> rows_dict, cols_dict, edges_dict;
+  phmap::flat_hash_map<rel_t, vector<int64_t>> rows_dict, cols_dict, edges_dict;
   for (const auto &kv : colptr_dict) {
     const auto &rel_type = kv.key();
     rows_dict[rel_type];
@@ -156,41 +174,82 @@ hetero_sample(const vector<node_t> &node_types,
     const torch::Tensor &input_node = kv.value();
     const auto *input_node_data = input_node.data_ptr<int64_t>();
 
+    int64_t *node_time_data;
+    if (temporal) {
+      const torch::Tensor &node_time = node_time_dict.at(node_type);
+      node_time_data = node_time.data_ptr<int64_t>();
+    }
+
     auto &samples = samples_dict.at(node_type);
+    auto &temp_samples = temp_samples_dict.at(node_type);
     auto &to_local_node = to_local_node_dict.at(node_type);
+    auto &temp_to_local_node = temp_to_local_node_dict.at(node_type);
+    auto &root_time = root_time_dict.at(node_type);
     for (int64_t i = 0; i < input_node.numel(); i++) {
       const auto &v = input_node_data[i];
+      if (temporal) {
+        temp_samples.push_back({v, i});
+        temp_to_local_node.insert({{v, i}, i});
+      } else {
         samples.push_back(v);
         to_local_node.insert({v, i});
       }
+      if (temporal)
+        root_time.push_back(node_time_data[v]);
+    }
   }
 
-  unordered_map<node_t, pair<int64_t, int64_t>> slice_dict;
+  phmap::flat_hash_map<node_t, pair<int64_t, int64_t>> slice_dict;
+  if (temporal) {
+    for (const auto &kv : temp_samples_dict) {
+      slice_dict[kv.first] = {0, kv.second.size()};
+    }
+  } else {
     for (const auto &kv : samples_dict)
       slice_dict[kv.first] = {0, kv.second.size()};
+  }
 
-  for (int64_t ell = 0; ell < num_hops; ell++) {
+  vector<rel_t> all_rel_types;
   for (const auto &kv : num_neighbors_dict) {
-      const auto &rel_type = kv.key();
+    all_rel_types.push_back(kv.key());
+  }
+  std::sort(all_rel_types.begin(), all_rel_types.end());
+
+  for (int64_t ell = 0; ell < num_hops; ell++) {
+    for (const auto &rel_type : all_rel_types) {
       const auto &edge_type = to_edge_type[rel_type];
       const auto &src_node_type = get<0>(edge_type);
       const auto &dst_node_type = get<2>(edge_type);
-      const auto num_samples = kv.value()[ell];
+      const auto num_samples = num_neighbors_dict.at(rel_type)[ell];
       const auto &dst_samples = samples_dict.at(dst_node_type);
+      const auto &temp_dst_samples = temp_samples_dict.at(dst_node_type);
       auto &src_samples = samples_dict.at(src_node_type);
+      auto &temp_src_samples = temp_samples_dict.at(src_node_type);
       auto &to_local_src_node = to_local_node_dict.at(src_node_type);
+      auto &temp_to_local_src_node = temp_to_local_node_dict.at(src_node_type);
 
-      const auto *colptr_data = ((torch::Tensor)colptr_dict.at(rel_type)).data_ptr<int64_t>();
-      const auto *row_data = ((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>();
+      const torch::Tensor &colptr = colptr_dict.at(rel_type);
+      const auto *colptr_data = colptr.data_ptr<int64_t>();
+      const torch::Tensor &row = row_dict.at(rel_type);
+      const auto *row_data = row.data_ptr<int64_t>();
 
       auto &rows = rows_dict.at(rel_type);
       auto &cols = cols_dict.at(rel_type);
       auto &edges = edges_dict.at(rel_type);
 
+      // For temporal sampling, sampled nodes cannot have a timestamp greater
+      // than the timestamp of the root nodes:
+      const auto &dst_root_time = root_time_dict.at(dst_node_type);
+      auto &src_root_time = root_time_dict.at(src_node_type);
+
       const auto &begin = slice_dict.at(dst_node_type).first;
       const auto &end = slice_dict.at(dst_node_type).second;
       for (int64_t i = begin; i < end; i++) {
-        const auto &w = dst_samples[i];
+        const auto &w = temporal ? temp_dst_samples[i].first : dst_samples[i];
+        const int64_t root_w = temporal ? temp_dst_samples[i].second : -1;
+        int64_t dst_time = 0;
+        if (temporal)
+          dst_time = dst_root_time[i];
         const auto &col_start = colptr_data[w];
         const auto &col_end = colptr_data[w + 1];
         const auto col_count = col_end - col_start;
@@ -199,8 +258,26 @@ hetero_sample(const vector<node_t> &node_types,
           continue;
 
         if ((num_samples < 0) || (!replace && (num_samples >= col_count))) {
+          // Select all neighbors:
           for (int64_t offset = col_start; offset < col_end; offset++) {
             const int64_t &v = row_data[offset];
+            if (temporal) {
+              if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
+                continue;
+              // force disjoint of computation tree based on source batch idx.
+              // note that the sampling always needs to have directed=True
+              // for temporal case
+              // to_local_src_node is not used for temporal / directed case
+              const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
+              if (res.second) {
+                temp_src_samples.push_back({v, root_w});
+                src_root_time.push_back(dst_time);
+              }
+
+              cols.push_back(i);
+              rows.push_back(res.first->second);
+              edges.push_back(offset);
+            } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
               if (res.second)
                 src_samples.push_back(v);
@@ -210,10 +287,30 @@ hetero_sample(const vector<node_t> &node_types,
                 edges.push_back(offset);
               }
             }
+          }
         } else if (replace) {
-          for (int64_t j = 0; j < num_samples; j++) {
-            const int64_t offset = col_start + rand() % col_count;
+          // Sample with replacement:
+          int64_t num_neighbors = 0;
+          while (num_neighbors < num_samples) {
+            const int64_t offset = col_start + uniform_randint(col_count);
             const int64_t &v = row_data[offset];
+            if (temporal) {
+              // TODO Infinity loop if no neighbor satisfies time constraint:
+              if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
+                continue;
+              // force disjoint of computation tree based on source batch idx.
+              // note that the sampling always needs to have directed=True
+              // for temporal case
+              const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
+              if (res.second) {
+                temp_src_samples.push_back({v, root_w});
+                src_root_time.push_back(dst_time);
+              }
+
+              cols.push_back(i);
+              rows.push_back(res.first->second);
+              edges.push_back(offset);
+            } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
               if (res.second)
                 src_samples.push_back(v);
@@ -223,16 +320,35 @@ hetero_sample(const vector<node_t> &node_types,
                 edges.push_back(offset);
               }
             }
+            num_neighbors += 1;
+          }
         } else {
+          // Sample without replacement:
           unordered_set<int64_t> rnd_indices;
           for (int64_t j = col_count - num_samples; j < col_count; j++) {
-            int64_t rnd = rand() % j;
+            int64_t rnd = uniform_randint(j);
             if (!rnd_indices.insert(rnd).second) {
               rnd = j;
               rnd_indices.insert(j);
             }
             const int64_t offset = col_start + rnd;
             const int64_t &v = row_data[offset];
+            if (temporal) {
+              if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
+                continue;
+              // force disjoint of computation tree based on source batch idx.
+              // note that the sampling always needs to have directed=True
+              // for temporal case
+              const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
+              if (res.second) {
+                temp_src_samples.push_back({v, root_w});
+                src_root_time.push_back(dst_time);
+              }
+
+              cols.push_back(i);
+              rows.push_back(res.first->second);
+              edges.push_back(offset);
+            } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
               if (res.second)
                 src_samples.push_back(v);
@@ -245,14 +361,22 @@ hetero_sample(const vector<node_t> &node_types,
           }
         }
       }
+    }
 
-    for (const auto &kv : samples_dict) {
-      slice_dict[kv.first] = {slice_dict.at(kv.first).second, kv.second.size()};
+    if (temporal) {
+      for (const auto &kv : temp_samples_dict) {
+        slice_dict[kv.first] = {0, kv.second.size()};
+      }
+    } else {
+      for (const auto &kv : samples_dict)
+        slice_dict[kv.first] = {0, kv.second.size()};
     }
   }
 
+  // Temporal sample disable undirected
+  assert(!(temporal && !directed));
   if (!directed) { // Construct the subgraph among the sampled nodes:
-    unordered_map<int64_t, int64_t>::iterator iter;
+    phmap::flat_hash_map<int64_t, int64_t>::iterator iter;
     for (const auto &kv : colptr_dict) {
       const auto &rel_type = kv.key();
       const auto &edge_type = to_edge_type[rel_type];
@@ -262,7 +386,8 @@ hetero_sample(const vector<node_t> &node_types,
       auto &to_local_src_node = to_local_node_dict.at(src_node_type);
 
       const auto *colptr_data = ((torch::Tensor)kv.value()).data_ptr<int64_t>();
-      const auto *row_data = ((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>();
+      const auto *row_data =
+          ((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>();
 
       auto &rows = rows_dict.at(rel_type);
       auto &cols = cols_dict.at(rel_type);
@@ -285,6 +410,18 @@ hetero_sample(const vector<node_t> &node_types,
     }
   }
 
+  // Construct samples dictionary from temporal sample dictionary.
+  if (temporal) {
+    for (const auto &kv : temp_samples_dict) {
+      const auto &node_type = kv.first;
+      const auto &samples = kv.second;
+      samples_dict[node_type].reserve(samples.size());
+      for (const auto &v : samples) {
+        samples_dict[node_type].push_back(v.first);
+      }
+    }
+  }
+
   return make_tuple(from_vector<node_t, int64_t>(samples_dict),
                     from_vector<rel_t, int64_t>(rows_dict),
                     from_vector<rel_t, int64_t>(cols_dict),
@@ -320,21 +457,51 @@ hetero_neighbor_sample_cpu(
     const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
     const int64_t num_hops, const bool replace, const bool directed) {
 
+  c10::Dict<node_t, torch::Tensor> node_time_dict; // Empty dictionary.
+
   if (replace && directed) {
-    return hetero_sample<true, true>(node_types, edge_types, colptr_dict,
-                                     row_dict, input_node_dict,
-                                     num_neighbors_dict, num_hops);
+    return hetero_sample<true, true, false>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
   } else if (replace && !directed) {
-    return hetero_sample<true, false>(node_types, edge_types, colptr_dict,
-                                      row_dict, input_node_dict,
-                                      num_neighbors_dict, num_hops);
+    return hetero_sample<true, false, false>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
   } else if (!replace && directed) {
-    return hetero_sample<false, true>(node_types, edge_types, colptr_dict,
-                                      row_dict, input_node_dict,
-                                      num_neighbors_dict, num_hops);
+    return hetero_sample<false, true, false>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
+  } else {
+    return hetero_sample<false, false, false>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
+  }
+}
+
+tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
+      c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
+hetero_temporal_neighbor_sample_cpu(
+    const vector<node_t> &node_types, const vector<edge_t> &edge_types,
+    const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
+    const c10::Dict<rel_t, torch::Tensor> &row_dict,
+    const c10::Dict<node_t, torch::Tensor> &input_node_dict,
+    const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
+    const c10::Dict<node_t, torch::Tensor> &node_time_dict,
+    const int64_t num_hops, const bool replace, const bool directed) {
+  AT_ASSERTM(directed, "Temporal sampling requires 'directed' sampling");
+  if (replace) {
+    // We assume that directed = True for temporal sampling
+    // The current implementation uses disjoint computation trees
+    // to tackle the case of the same node sampled having different
+    // root time constraint.
+    // In future, we could extend to directed = False case,
+    // allowing additional edges within each computation tree.
+    return hetero_sample<true, true, true>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
   } else {
-    return hetero_sample<false, false>(node_types, edge_types, colptr_dict,
-                                       row_dict, input_node_dict,
-                                       num_neighbors_dict, num_hops);
+    return hetero_sample<false, true, true>(
+        node_types, edge_types, colptr_dict, row_dict, input_node_dict,
+        num_neighbors_dict, node_time_dict, num_hops);
   }
 }
\ No newline at end of file

csrc/cpu/neighbor_sample_cpu.h

@@ -22,3 +22,15 @@ hetero_neighbor_sample_cpu(
     const c10::Dict<node_t, torch::Tensor> &input_node_dict,
     const c10::Dict<rel_t, std::vector<int64_t>> &num_neighbors_dict,
     const int64_t num_hops, const bool replace, const bool directed);
+
+std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
+           c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
+hetero_temporal_neighbor_sample_cpu(
+    const std::vector<node_t> &node_types,
+    const std::vector<edge_t> &edge_types,
+    const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
+    const c10::Dict<rel_t, torch::Tensor> &row_dict,
+    const c10::Dict<node_t, torch::Tensor> &input_node_dict,
+    const c10::Dict<rel_t, std::vector<int64_t>> &num_neighbors_dict,
+    const c10::Dict<node_t, torch::Tensor> &node_time_dict,
+    const int64_t num_hops, const bool replace, const bool directed);

csrc/cpu/relabel_cpu.cpp

@@ -64,7 +64,7 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
   std::unordered_map<int64_t, int64_t> n_id_map;
   std::unordered_map<int64_t, int64_t>::iterator it;
 
-  auto out_rowptr = torch::empty(idx.numel() + 1, rowptr.options());
+  auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options());
   auto out_rowptr_data = out_rowptr.data_ptr<int64_t>();
 
   out_rowptr_data[0] = 0;
@@ -76,12 +76,12 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
     out_rowptr_data[i + 1] = offset;
   }
 
-  auto out_col = torch::empty(offset, col.options());
+  auto out_col = torch::empty({offset}, col.options());
   auto out_col_data = out_col.data_ptr<int64_t>();
 
   torch::optional<torch::Tensor> out_value = torch::nullopt;
   if (optional_value.has_value()) {
-    out_value = torch::empty(offset, optional_value.value().options());
+    out_value = torch::empty({offset}, optional_value.value().options());
 
     AT_DISPATCH_ALL_TYPES(optional_value.value().scalar_type(), "relabel", [&] {
       auto value_data = optional_value.value().data_ptr<scalar_t>();

csrc/cpu/sample_cpu.cpp

@@ -15,13 +15,11 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
   CHECK_CPU(idx);
   CHECK_INPUT(idx.dim() == 1);
 
-  srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
-
   auto rowptr_data = rowptr.data_ptr<int64_t>();
   auto col_data = col.data_ptr<int64_t>();
   auto idx_data = idx.data_ptr<int64_t>();
 
-  auto out_rowptr = torch::empty(idx.numel() + 1, rowptr.options());
+  auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options());
   auto out_rowptr_data = out_rowptr.data_ptr<int64_t>();
   out_rowptr_data[0] = 0;
 
@@ -69,7 +67,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
 
       if (row_count > 0) {
         for (int64_t j = 0; j < num_neighbors; j++) {
-          e = row_start + rand() % row_count;
+          e = row_start + uniform_randint(row_count);
           c = col_data[e];
 
           if (n_id_map.count(c) == 0) {
@@ -96,7 +94,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
       } else { // See: https://www.nowherenearithaca.com/2013/05/
                //      robert-floyds-tiny-and-beautiful.html
         for (int64_t j = row_count - num_neighbors; j < row_count; j++) {
-          if (!perm.insert(rand() % j).second)
+          if (!perm.insert(uniform_randint(j)).second)
             perm.insert(j);
         }
       }
@@ -119,9 +117,9 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
   auto out_n_id = torch::from_blob(n_ids.data(), {N}, col.options()).clone();
 
   int64_t E = out_rowptr_data[idx.numel()];
-  auto out_col = torch::empty(E, col.options());
+  auto out_col = torch::empty({E}, col.options());
   auto out_col_data = out_col.data_ptr<int64_t>();
-  auto out_e_id = torch::empty(E, col.options());
+  auto out_e_id = torch::empty({E}, col.options());
   auto out_e_id_data = out_e_id.data_ptr<int64_t>();
 
   i = 0;