pax_global_header00006660000000000000000000000064145076146610014524gustar00rootroot0000000000000052 comment=7d22892fb47626d2c5b0171769d44bdc8edd606c pytorch_sparse-0.6.18/000077500000000000000000000000001450761466100146655ustar00rootroot00000000000000pytorch_sparse-0.6.18/.coveragerc000066400000000000000000000001621450761466100170050ustar00rootroot00000000000000[run] source=torch_sparse [report] exclude_lines = pragma: no cover torch.jit.script raise except pytorch_sparse-0.6.18/.github/000077500000000000000000000000001450761466100162255ustar00rootroot00000000000000pytorch_sparse-0.6.18/.github/workflows/000077500000000000000000000000001450761466100202625ustar00rootroot00000000000000pytorch_sparse-0.6.18/.github/workflows/building-conda.yml000066400000000000000000000061141450761466100236660ustar00rootroot00000000000000name: Building Conda on: [workflow_dispatch] jobs: conda-build: runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: # We have trouble building for Windows - drop for now. os: [ubuntu-20.04, macos-11] # windows-2019 python-version: ['3.8', '3.9', '3.10', '3.11'] torch-version: [2.0.0, 2.1.0] cuda-version: ['cpu', 'cu117', 'cu118', 'cu121'] exclude: - torch-version: 2.0.0 cuda-version: 'cu121' - torch-version: 2.1.0 cuda-version: 'cu117' - os: macos-11 cuda-version: 'cu117' - os: macos-11 cuda-version: 'cu118' - os: macos-11 cuda-version: 'cu121' steps: - uses: actions/checkout@v2 with: submodules: 'recursive' - name: Set up Conda for Python ${{ matrix.python-version }} uses: conda-incubator/setup-miniconda@v2 with: python-version: ${{ matrix.python-version }} - name: Free Disk Space (Ubuntu) if: ${{ runner.os == 'Linux' }} uses: jlumbroso/free-disk-space@main - name: Install Conda packages run: | conda install conda-build conda-verify --yes shell: bash -l {0} - name: Install CUDA ${{ matrix.cuda-version }} if: ${{ matrix.cuda-version != 'cpu' }} run: | bash .github/workflows/cuda/${{ matrix.cuda-version }}-${{ runner.os }}.sh shell: bash - name: Install METIS if: ${{ runner.os != 'Windows' }} run: | bash .github/workflows/metis.sh - name: Install METIS on Windows if: ${{ runner.os == 'Windows' }} run: | bash .github/workflows/metis-${{ runner.os }}.sh shell: bash -l {0} - name: Build Conda package for CPU if: ${{ matrix.cuda-version == 'cpu' }} run: | FORCE_CUDA=0 TORCH_CUDA_ARCH_LIST=0 ./conda/pytorch-sparse/build_conda.sh ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }} shell: bash -l {0} - name: Build Conda package for GPU if: ${{ matrix.cuda-version != 'cpu' }} run: | source .github/workflows/cuda/${{ matrix.cuda-version }}-${{ runner.os }}-env.sh ./conda/pytorch-sparse/build_conda.sh ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }} shell: bash -l {0} - name: Publish Conda package on organization channel run: | conda install anaconda-client --yes anaconda upload --force --label main $HOME/conda-bld/*/*.tar.bz2 env: ANACONDA_API_TOKEN: ${{ secrets.PYG_ANACONDA_TOKEN }} shell: bash -l {0} - name: Publish Conda package on personal channel run: | conda install anaconda-client --yes anaconda upload --force --label main $HOME/conda-bld/*/*.tar.bz2 env: ANACONDA_API_TOKEN: ${{ secrets.RUSTY1S_ANACONDA_TOKEN }} shell: bash -l {0} pytorch_sparse-0.6.18/.github/workflows/building.yml000066400000000000000000000077371450761466100226200ustar00rootroot00000000000000name: Building Wheels on: [workflow_dispatch] jobs: wheel: runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: os: [ubuntu-20.04, macos-11, windows-2019] python-version: ['3.8', '3.9', '3.10', '3.11'] torch-version: [2.0.0, 2.1.0] cuda-version: ['cpu', 'cu117', 'cu118', 'cu121'] exclude: - torch-version: 2.0.0 cuda-version: 'cu121' - torch-version: 2.1.0 cuda-version: 'cu117' - os: macos-11 cuda-version: 'cu117' - os: macos-11 cuda-version: 'cu118' - os: macos-11 cuda-version: 'cu121' steps: - uses: actions/checkout@v2 with: submodules: 'recursive' - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v2 with: python-version: ${{ matrix.python-version }} - name: Upgrade pip run: | pip install --upgrade setuptools pip install scipy==1.10.1 # Python 3.8 support - name: Free Disk Space (Ubuntu) if: ${{ runner.os == 'Linux' }} uses: jlumbroso/free-disk-space@main - name: Install CUDA ${{ matrix.cuda-version }} if: ${{ matrix.cuda-version != 'cpu' }} run: | bash .github/workflows/cuda/${{ matrix.cuda-version }}-${{ runner.os }}.sh - name: Install PyTorch ${{ matrix.torch-version }}+${{ matrix.cuda-version }} run: | pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/${{ matrix.cuda-version }} python -c "import torch; print('PyTorch:', torch.__version__)" python -c "import torch; print('CUDA:', torch.version.cuda)" - name: Set version if: ${{ runner.os != 'macOS' }} run: | VERSION=`sed -n "s/^__version__ = '\(.*\)'/\1/p" torch_sparse/__init__.py` TORCH_VERSION=`echo "pt${{ matrix.torch-version }}" | sed "s/..$//" | sed "s/\.//g"` CUDA_VERSION=`echo ${{ matrix.cuda-version }}` echo "New version name: $VERSION+$TORCH_VERSION$CUDA_VERSION" sed -i "s/$VERSION/$VERSION+$TORCH_VERSION$CUDA_VERSION/" setup.py sed -i "s/$VERSION/$VERSION+$TORCH_VERSION$CUDA_VERSION/" torch_sparse/__init__.py shell: bash - name: Install torch-scatter run: | pip install torch-scatter -f https://data.pyg.org/whl/torch-${{ matrix.torch-version }}+${{ matrix.cuda-version }}.html - name: Install METIS if: ${{ runner.os != 'Windows' }} run: | bash .github/workflows/metis.sh - name: Install METIS on Windows if: ${{ runner.os == 'Windows' }} run: | bash .github/workflows/metis-${{ runner.os }}.sh - name: Install main package for CPU if: ${{ matrix.cuda-version == 'cpu' }} run: | FORCE_ONLY_CPU=1 WITH_METIS=1 python setup.py develop shell: bash - name: Install main package for GPU if: ${{ matrix.cuda-version != 'cpu' }} run: | source .github/workflows/cuda/${{ matrix.cuda-version }}-${{ runner.os }}-env.sh WITH_METIS=1 python setup.py develop shell: bash - name: Test installation run: | python -c "import torch_sparse; print('torch-sparse:', torch_sparse.__version__)" - name: Build wheel run: | pip install wheel python setup.py bdist_wheel --dist-dir=dist - name: Configure AWS uses: aws-actions/configure-aws-credentials@v1 with: aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} aws-region: us-west-1 - name: Upload wheel run: | aws s3 sync dist s3://data.pyg.org/whl/torch-${{ matrix.torch-version }}+${{ matrix.cuda-version }} --grants read=uri=http://acs.amazonaws.com/groups/global/AllUsers pytorch_sparse-0.6.18/.github/workflows/cuda/000077500000000000000000000000001450761466100211765ustar00rootroot00000000000000pytorch_sparse-0.6.18/.github/workflows/cuda/cu101-Linux-env.sh000066400000000000000000000003131450761466100242430ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-10.1 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" pytorch_sparse-0.6.18/.github/workflows/cuda/cu101-Linux.sh000077500000000000000000000013361450761466100234660ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda-repo-${OS}-10-1-local-10.1.243-418.87.00_1.0-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-10-1-local-10.1.243-418.87.00_1.0-1_amd64.deb sudo apt-key add /var/cuda-repo-10-1-local-10.1.243-418.87.00/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-10-1 cuda-libraries-dev-10-1 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda-repo-${OS}-10-1-local-10.1.243-418.87.00_1.0-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu101-Windows-env.sh000066400000000000000000000004341450761466100246020ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v10.1 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" pytorch_sparse-0.6.18/.github/workflows/cuda/cu101-Windows.sh000077500000000000000000000020761450761466100240230ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=10.1 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}/Prod/local_installers/ export CUDA_FILE=cuda_${CUDA_SHORT}.243_426.00_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu102-Linux-env.sh000066400000000000000000000003131450761466100242440ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-10.2 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" pytorch_sparse-0.6.18/.github/workflows/cuda/cu102-Linux.sh000077500000000000000000000013321450761466100234630ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda-repo-${OS}-10-2-local-10.2.89-440.33.01_1.0-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-10-2-local-10.2.89-440.33.01_1.0-1_amd64.deb sudo apt-key add /var/cuda-repo-10-2-local-10.2.89-440.33.01/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-10-2 cuda-libraries-dev-10-2 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda-repo-${OS}-10-2-local-10.2.89-440.33.01_1.0-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu102-Windows-env.sh000066400000000000000000000004341450761466100246030ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v10.2 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" pytorch_sparse-0.6.18/.github/workflows/cuda/cu102-Windows.sh000077500000000000000000000020741450761466100240220ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=10.2 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}/Prod/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.89_441.22_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu111-Linux-env.sh000066400000000000000000000003231450761466100242450ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.1 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu111-Linux.sh000077500000000000000000000012711450761466100234650ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda-repo-${OS}-11-1-local_11.1.1-455.32.00-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-1-local_11.1.1-455.32.00-1_amd64.deb sudo apt-key add /var/cuda-repo-${OS}-11-1-local/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-11-1 cuda-libraries-dev-11-1 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda-repo-${OS}-11-1-local_11.1.1-455.32.00-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu111-Windows-env.sh000066400000000000000000000004441450761466100246040ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.1 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu111-Windows.sh000077500000000000000000000020701450761466100240160ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.1 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.1/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.1_456.81_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu113-Linux-env.sh000066400000000000000000000003231450761466100242470ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.3 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu113-Linux.sh000077500000000000000000000012711450761466100234670ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb sudo apt-key add /var/cuda-repo-${OS}-11-3-local/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-11-3 cuda-libraries-dev-11-3 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu113-Windows-env.sh000066400000000000000000000004441450761466100246060ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.3 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu113-Windows.sh000077500000000000000000000021151450761466100240200ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.3 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.0_465.89_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu115-Linux-env.sh000066400000000000000000000003231450761466100242510ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.5 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu115-Linux.sh000077500000000000000000000012711450761466100234710ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.5.2/local_installers/cuda-repo-${OS}-11-5-local_11.5.2-495.29.05-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-5-local_11.5.2-495.29.05-1_amd64.deb sudo apt-key add /var/cuda-repo-${OS}-11-5-local/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-11-5 cuda-libraries-dev-11-5 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.5.2/local_installers/cuda-repo-${OS}-11-5-local_11.5.2-495.29.05-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu115-Windows-env.sh000066400000000000000000000004141450761466100246050ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.3 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="6.0+PTX" pytorch_sparse-0.6.18/.github/workflows/cuda/cu115-Windows.sh000077500000000000000000000022221450761466100240210ustar00rootroot00000000000000#!/bin/bash # TODO We currently use CUDA 11.3 to build CUDA 11.5 Windows wheels # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.3 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.0_465.89_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu116-Linux-env.sh000066400000000000000000000003231450761466100242520ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.6 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu116-Linux.sh000077500000000000000000000012711450761466100234720ustar00rootroot00000000000000#!/bin/bash OS=ubuntu1804 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda-repo-${OS}-11-6-local_11.6.2-510.47.03-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-6-local_11.6.2-510.47.03-1_amd64.deb sudo apt-key add /var/cuda-repo-${OS}-11-6-local/7fa2af80.pub sudo apt-get -qq update sudo apt install cuda-nvcc-11-6 cuda-libraries-dev-11-6 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.5.2/local_installers/cuda-repo-${OS}-11-6-local_11.6.2-510.47.03-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu116-Windows-env.sh000066400000000000000000000004141450761466100246060ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.3 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="6.0+PTX" pytorch_sparse-0.6.18/.github/workflows/cuda/cu116-Windows.sh000077500000000000000000000022221450761466100240220ustar00rootroot00000000000000#!/bin/bash # TODO We currently use CUDA 11.3 to build CUDA 11.6 Windows wheels # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.3 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.0_465.89_win10.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu117-Linux-env.sh000066400000000000000000000003231450761466100242530ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.7 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu117-Linux.sh000077500000000000000000000013131450761466100234700ustar00rootroot00000000000000#!/bin/bash OS=ubuntu2004 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-${OS}-11-7-local_11.7.1-515.65.01-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-7-local_11.7.1-515.65.01-1_amd64.deb sudo cp /var/cuda-repo-${OS}-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/ sudo apt-get -qq update sudo apt install cuda-nvcc-11-7 cuda-libraries-dev-11-7 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-${OS}-11-7-local_11.7.1-515.65.01-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu117-Windows-env.sh000066400000000000000000000004141450761466100246070ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.7 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="6.0+PTX" pytorch_sparse-0.6.18/.github/workflows/cuda/cu117-Windows.sh000077500000000000000000000021171450761466100240260ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.7 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.1/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.1_516.94_windows.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu118-Linux-env.sh000066400000000000000000000003231450761466100242540ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-11.8 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" pytorch_sparse-0.6.18/.github/workflows/cuda/cu118-Linux.sh000077500000000000000000000013131450761466100234710ustar00rootroot00000000000000#!/bin/bash OS=ubuntu2004 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-${OS}-11-8-local_11.8.0-520.61.05-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-11-8-local_11.8.0-520.61.05-1_amd64.deb sudo cp /var/cuda-repo-${OS}-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/ sudo apt-get -qq update sudo apt install cuda-nvcc-11-8 cuda-libraries-dev-11-8 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-${OS}-11-8-local_11.8.0-520.61.05-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu118-Windows-env.sh000066400000000000000000000004141450761466100246100ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.8 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="6.0+PTX" pytorch_sparse-0.6.18/.github/workflows/cuda/cu118-Windows.sh000077500000000000000000000021171450761466100240270ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=11.8 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.0_522.06_windows.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/cuda/cu121-Linux-env.sh000066400000000000000000000003231450761466100242460ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/usr/local/cuda-12.1 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} PATH=${CUDA_HOME}/bin:${PATH} export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0" pytorch_sparse-0.6.18/.github/workflows/cuda/cu121-Linux.sh000077500000000000000000000013131450761466100234630ustar00rootroot00000000000000#!/bin/bash OS=ubuntu2004 wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget -nv https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda-repo-${OS}-12-1-local_12.1.1-530.30.02-1_amd64.deb sudo dpkg -i cuda-repo-${OS}-12-1-local_12.1.1-530.30.02-1_amd64.deb sudo cp /var/cuda-repo-${OS}-12-1-local/cuda-*-keyring.gpg /usr/share/keyrings/ sudo apt-get -qq update sudo apt install cuda-nvcc-12-1 cuda-libraries-dev-12-1 sudo apt clean rm -f https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda-repo-${OS}-12-1-local_12.1.1-530.30.02-1_amd64.deb pytorch_sparse-0.6.18/.github/workflows/cuda/cu121-Windows-env.sh000066400000000000000000000004141450761466100246020ustar00rootroot00000000000000#!/bin/bash CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v12.1 PATH=${CUDA_HOME}/bin:$PATH PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="6.0+PTX" pytorch_sparse-0.6.18/.github/workflows/cuda/cu121-Windows.sh000077500000000000000000000021171450761466100240210ustar00rootroot00000000000000#!/bin/bash # Install NVIDIA drivers, see: # https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" export CUDA_SHORT=12.1 export CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.1/local_installers export CUDA_FILE=cuda_${CUDA_SHORT}.1_531.14_windows.exe # Install CUDA: curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}" echo "" echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" pytorch_sparse-0.6.18/.github/workflows/linting.yml000066400000000000000000000006371450761466100224570ustar00rootroot00000000000000name: Linting on: push: branches: - master pull_request: jobs: flake8: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 - name: Set up Python uses: actions/setup-python@v2 with: python-version: 3.8 - name: Install dependencies run: | pip install flake8 - name: Run linting run: | flake8 . pytorch_sparse-0.6.18/.github/workflows/metis-Windows.sh000077500000000000000000000015251450761466100233750ustar00rootroot00000000000000#!/bin/bash METIS=metis-5.1.0 curl -k -L "https://web.archive.org/web/20170712055800/http://glaros.dtc.umn.edu/gkhome/fetch/sw/metis/${METIS}.tar.gz" --output "${METIS}.tar.gz" tar -xvzf "${METIS}.tar.gz" rm -f "${METIS}.tar.gz" cd "${METIS}" || exit sed -i.bak -e 's/IDXTYPEWIDTH 32/IDXTYPEWIDTH 64/g' include/metis.h # Fix GKlib on Windows: https://github.com/jlblancoc/suitesparse-metis-for-windows/issues/6 sed -i.bak -e '61,69d' GKlib/gk_arch.h cd build || exit cmake .. -A x64 # Ensure we are building with x64 cmake --build . --config "Release" --target ALL_BUILD cp libmetis/Release/metis.lib /c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2019/Enterprise/VC/Tools/MSVC/14.29.30133/lib/x64 cp ../include/metis.h /c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2019/Enterprise/VC/Tools/MSVC/14.29.30133/include rm -f "${METIS}" pytorch_sparse-0.6.18/.github/workflows/metis.sh000077500000000000000000000006701450761466100217450ustar00rootroot00000000000000#!/bin/bash METIS=metis-5.1.0 wget -nv "https://web.archive.org/web/20211119110155/http://glaros.dtc.umn.edu/gkhome/fetch/sw/metis/${METIS}.tar.gz" tar -xvzf "${METIS}.tar.gz" rm -f "${METIS}.tar.gz" cd "${METIS}" || exit sed -i.bak -e 's/IDXTYPEWIDTH 32/IDXTYPEWIDTH 64/g' include/metis.h make config make sudo make install sudo cp /usr/local/include/metis.h /usr/include/ sudo cp /usr/local/lib/libmetis.a /usr/lib/ rm -f "${METIS}" pytorch_sparse-0.6.18/.github/workflows/stale.yml000066400000000000000000000014261450761466100221200ustar00rootroot00000000000000name: "Close stale issues and PRs" on: schedule: # Every day at 00:00 - cron: "0 0 * * *" workflow_dispatch: jobs: stale: runs-on: ubuntu-latest steps: - uses: actions/stale@v4.0.0 with: stale-issue-message: 'This issue had no activity for **6 months**. It will be closed in **2 weeks** unless there is some new activity. Is this issue already resolved?' stale-issue-label: 'stale' exempt-issue-labels: 'bug,enhancement,good first issue' stale-pr-message: 'This pull request had no activity for **6 months**. It will be closed in **2 weeks** unless there is some new activity.' stale-pr-label: 'stale' days-before-stale: 180 days-before-close: 14 operations-per-run: 200 pytorch_sparse-0.6.18/.github/workflows/testing.yml000066400000000000000000000031201450761466100224560ustar00rootroot00000000000000name: Testing on: push: branches: - master pull_request: jobs: pytest: runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: os: [ubuntu-latest, windows-2019] python-version: [3.8] torch-version: [2.0.0, 2.1.0] steps: - uses: actions/checkout@v2 with: submodules: 'recursive' - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v2 with: python-version: ${{ matrix.python-version }} - name: Install PyTorch ${{ matrix.torch-version }} run: | pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/cpu - name: Install torch-scatter run: | pip install torch-scatter -f https://data.pyg.org/whl/torch-${{ matrix.torch-version }}+cpu.html - name: Install METIS if: ${{ runner.os != 'Windows' }} run: | bash .github/workflows/metis.sh - name: Install METIS on Windows if: ${{ runner.os == 'Windows' }} run: | bash .github/workflows/metis-${{ runner.os }}.sh - name: Install main package run: | pip install scipy==1.10.1 # Python 3.8 support python setup.py develop env: WITH_METIS: 1 - name: Run test-suite run: | pip install pytest pytest-cov pytest --cov --cov-report=xml - name: Upload coverage uses: codecov/codecov-action@v1 if: success() with: fail_ci_if_error: false pytorch_sparse-0.6.18/.gitignore000066400000000000000000000001221450761466100166500ustar00rootroot00000000000000__pycache__/ build/ dist/ alpha/ .cache/ .eggs/ *.egg-info/ .coverage *.so .idea/ pytorch_sparse-0.6.18/.gitmodules000066400000000000000000000002121450761466100170350ustar00rootroot00000000000000[submodule "third_party/parallel-hashmap"] path = third_party/parallel-hashmap url = https://github.com/greg7mdp/parallel-hashmap.git pytorch_sparse-0.6.18/CMakeLists.txt000066400000000000000000000070261450761466100174320ustar00rootroot00000000000000cmake_minimum_required(VERSION 3.10) project(torchsparse) set(CMAKE_CXX_STANDARD 14) set(TORCHSPARSE_VERSION 0.6.18) set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/cmake) option(WITH_CUDA "Enable CUDA support" OFF) option(WITH_PYTHON "Link to Python when building" ON) option(WITH_METIS "Enable METIS support" OFF) 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) if (WITH_METIS) add_definitions(-DWITH_METIS) find_package(METIS) endif() 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() if (WITH_METIS) target_include_directories(${PROJECT_NAME} PRIVATE ${METIS_INCLUDE_DIRS}) target_link_libraries(${PROJECT_NAME} PRIVATE ${METIS_LIBRARIES}) endif() find_package(OpenMP) if (OPENMP_FOUND) set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") # set (CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler=${OpenMP_CXX_FLAGS}") set (CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") endif() set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchSparse) target_include_directories(${PROJECT_NAME} INTERFACE "$" $) 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 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 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() pytorch_sparse-0.6.18/LICENSE000066400000000000000000000020761450761466100156770ustar00rootroot00000000000000Copyright (c) 2020 Matthias Fey Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. pytorch_sparse-0.6.18/MANIFEST.in000066400000000000000000000006541450761466100164300ustar00rootroot00000000000000include README.md include LICENSE recursive-include csrc * recursive-include third_party * recursive-exclude third_party/parallel-hashmap/css * recursive-exclude third_party/parallel-hashmap/html * recursive-exclude third_party/parallel-hashmap/tests * recursive-exclude third_party/parallel-hashmap/examples * recursive-exclude third_party/parallel-hashmap/benchmark * recursive-exclude test * recursive-exclude benchmark * pytorch_sparse-0.6.18/README.md000066400000000000000000000240271450761466100161510ustar00rootroot00000000000000[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 # PyTorch Sparse [![PyPI Version][pypi-image]][pypi-url] [![Testing Status][testing-image]][testing-url] [![Linting Status][linting-image]][linting-url] [![Code Coverage][coverage-image]][coverage-url] -------------------------------------------------------------------------------- This package consists of a small extension library of optimized sparse matrix operations with autograd support. This package currently consists of the following methods: * **[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 ``` ### Binaries We alternatively provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://data.pyg.org/whl). #### PyTorch 2.1 To install the binaries for PyTorch 2.1.0, simply run ``` pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.1.0+${CUDA}.html ``` where `${CUDA}` should be replaced by either `cpu`, `cu118`, or `cu121` depending on your PyTorch installation. | | `cpu` | `cu118` | `cu121` | |-------------|-------|---------|---------| | **Linux** | ✅ | ✅ | ✅ | | **Windows** | ✅ | ✅ | ✅ | | **macOS** | ✅ | | | #### PyTorch 2.0 To install the binaries for PyTorch 2.0.0, simply run ``` pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html ``` where `${CUDA}` should be replaced by either `cpu`, `cu117`, or `cu118` depending on your PyTorch installation. | | `cpu` | `cu117` | `cu118` | |-------------|-------|---------|---------| | **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, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1 and PyTorch 1.13.0/1.13.1 (following the same procedure). For older versions, you need to explicitly specify the latest supported version number or install via `pip install --no-index` 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 ``` 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 from torch_sparse import coalesce 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 = coalesce(index, value, m=3, n=2) ``` ``` print(index) tensor([[0, 1, 1, 2], [1, 0, 1, 0]]) print(value) tensor([[6.0, 8.0], [7.0, 9.0], [3.0, 4.0], [5.0, 6.0]]) ``` ### Transpose ``` 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 ``` pytorch_sparse-0.6.18/benchmark/000077500000000000000000000000001450761466100166175ustar00rootroot00000000000000pytorch_sparse-0.6.18/benchmark/.gitignore000066400000000000000000000000141450761466100206020ustar00rootroot00000000000000*.mat *.tmp pytorch_sparse-0.6.18/benchmark/main.py000066400000000000000000000124571450761466100201260ustar00rootroot00000000000000import argparse import itertools import os.path as osp import time import torch import wget 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() elif torch.backends.mps.is_available(): import torch.mps torch.mps.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() elif torch.backends.mps.is_available(): import torch.mps torch.mps.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) pytorch_sparse-0.6.18/benchmark/ptr2ind.py000066400000000000000000000046011450761466100205540ustar00rootroot00000000000000import 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() pytorch_sparse-0.6.18/cmake/000077500000000000000000000000001450761466100157455ustar00rootroot00000000000000pytorch_sparse-0.6.18/cmake/FindMetis.cmake000066400000000000000000000214111450761466100206300ustar00rootroot00000000000000### # # @copyright (c) 2009-2014 The University of Tennessee and The University # of Tennessee Research Foundation. # All rights reserved. # @copyright (c) 2012-2014 Inria. All rights reserved. # @copyright (c) 2012-2014 Bordeaux INP, CNRS (LaBRI UMR 5800), Inria, Univ. Bordeaux. All rights reserved. # ### # # - Find METIS include dirs and libraries # Use this module by invoking find_package with the form: # find_package(METIS # [REQUIRED] # Fail with error if metis is not found # ) # # This module finds headers and metis library. # Results are reported in variables: # METIS_FOUND - True if headers and requested libraries were found # METIS_INCLUDE_DIRS - metis include directories # METIS_LIBRARY_DIRS - Link directories for metis libraries # METIS_LIBRARIES - metis component libraries to be linked # # The user can give specific paths where to find the libraries adding cmake # options at configure (ex: cmake path/to/project -DMETIS_DIR=path/to/metis): # METIS_DIR - Where to find the base directory of metis # METIS_INCDIR - Where to find the header files # METIS_LIBDIR - Where to find the library files # The module can also look for the following environment variables if paths # are not given as cmake variable: METIS_DIR, METIS_INCDIR, METIS_LIBDIR #============================================================================= # Copyright 2012-2013 Inria # Copyright 2012-2013 Emmanuel Agullo # Copyright 2012-2013 Mathieu Faverge # Copyright 2012 Cedric Castagnede # Copyright 2013 Florent Pruvost # # Distributed under the OSI-approved BSD License (the "License"); # see accompanying file MORSE-Copyright.txt for details. # # This software is distributed WITHOUT ANY WARRANTY; without even the # implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. # See the License for more information. #============================================================================= # (To distribute this file outside of Morse, substitute the full # License text for the above reference.) if (NOT METIS_FOUND) set(METIS_DIR "" CACHE PATH "Installation directory of METIS library") if (NOT METIS_FIND_QUIETLY) message(STATUS "A cache variable, namely METIS_DIR, has been set to specify the install directory of METIS") endif() endif() # Looking for include # ------------------- # Add system include paths to search include # ------------------------------------------ unset(_inc_env) set(ENV_METIS_DIR "$ENV{METIS_DIR}") set(ENV_METIS_INCDIR "$ENV{METIS_INCDIR}") if(ENV_METIS_INCDIR) list(APPEND _inc_env "${ENV_METIS_INCDIR}") elseif(ENV_METIS_DIR) list(APPEND _inc_env "${ENV_METIS_DIR}") list(APPEND _inc_env "${ENV_METIS_DIR}/include") list(APPEND _inc_env "${ENV_METIS_DIR}/include/metis") else() if(WIN32) string(REPLACE ":" ";" _inc_env "$ENV{INCLUDE}") else() string(REPLACE ":" ";" _path_env "$ENV{INCLUDE}") list(APPEND _inc_env "${_path_env}") string(REPLACE ":" ";" _path_env "$ENV{C_INCLUDE_PATH}") list(APPEND _inc_env "${_path_env}") string(REPLACE ":" ";" _path_env "$ENV{CPATH}") list(APPEND _inc_env "${_path_env}") string(REPLACE ":" ";" _path_env "$ENV{INCLUDE_PATH}") list(APPEND _inc_env "${_path_env}") endif() endif() list(APPEND _inc_env "${CMAKE_PLATFORM_IMPLICIT_INCLUDE_DIRECTORIES}") list(APPEND _inc_env "${CMAKE_C_IMPLICIT_INCLUDE_DIRECTORIES}") list(REMOVE_DUPLICATES _inc_env) # Try to find the metis header in the given paths # ------------------------------------------------- # call cmake macro to find the header path if(METIS_INCDIR) set(METIS_metis.h_DIRS "METIS_metis.h_DIRS-NOTFOUND") find_path(METIS_metis.h_DIRS NAMES metis.h HINTS ${METIS_INCDIR}) else() if(METIS_DIR) set(METIS_metis.h_DIRS "METIS_metis.h_DIRS-NOTFOUND") find_path(METIS_metis.h_DIRS NAMES metis.h HINTS ${METIS_DIR} PATH_SUFFIXES "include" "include/metis") else() set(METIS_metis.h_DIRS "METIS_metis.h_DIRS-NOTFOUND") find_path(METIS_metis.h_DIRS NAMES metis.h HINTS ${_inc_env}) endif() endif() mark_as_advanced(METIS_metis.h_DIRS) # If found, add path to cmake variable # ------------------------------------ if (METIS_metis.h_DIRS) set(METIS_INCLUDE_DIRS "${METIS_metis.h_DIRS}") else () set(METIS_INCLUDE_DIRS "METIS_INCLUDE_DIRS-NOTFOUND") if(NOT METIS_FIND_QUIETLY) message(STATUS "Looking for metis -- metis.h not found") endif() endif() # Looking for lib # --------------- # Add system library paths to search lib # -------------------------------------- unset(_lib_env) set(ENV_METIS_LIBDIR "$ENV{METIS_LIBDIR}") if(ENV_METIS_LIBDIR) list(APPEND _lib_env "${ENV_METIS_LIBDIR}") elseif(ENV_METIS_DIR) list(APPEND _lib_env "${ENV_METIS_DIR}") list(APPEND _lib_env "${ENV_METIS_DIR}/lib") else() if(WIN32) string(REPLACE ":" ";" _lib_env "$ENV{LIB}") else() if(APPLE) string(REPLACE ":" ";" _lib_env "$ENV{DYLD_LIBRARY_PATH}") else() string(REPLACE ":" ";" _lib_env "$ENV{LD_LIBRARY_PATH}") endif() list(APPEND _lib_env "${CMAKE_PLATFORM_IMPLICIT_LINK_DIRECTORIES}") list(APPEND _lib_env "${CMAKE_C_IMPLICIT_LINK_DIRECTORIES}") endif() endif() list(REMOVE_DUPLICATES _lib_env) # Try to find the metis lib in the given paths # ---------------------------------------------- # call cmake macro to find the lib path if(METIS_LIBDIR) set(METIS_metis_LIBRARY "METIS_metis_LIBRARY-NOTFOUND") find_library(METIS_metis_LIBRARY NAMES metis HINTS ${METIS_LIBDIR}) else() if(METIS_DIR) set(METIS_metis_LIBRARY "METIS_metis_LIBRARY-NOTFOUND") find_library(METIS_metis_LIBRARY NAMES metis HINTS ${METIS_DIR} PATH_SUFFIXES lib lib32 lib64) else() set(METIS_metis_LIBRARY "METIS_metis_LIBRARY-NOTFOUND") find_library(METIS_metis_LIBRARY NAMES metis HINTS ${_lib_env}) endif() endif() mark_as_advanced(METIS_metis_LIBRARY) # If found, add path to cmake variable # ------------------------------------ if (METIS_metis_LIBRARY) get_filename_component(metis_lib_path "${METIS_metis_LIBRARY}" PATH) # set cmake variables set(METIS_LIBRARIES "${METIS_metis_LIBRARY}") set(METIS_LIBRARY_DIRS "${metis_lib_path}") else () set(METIS_LIBRARIES "METIS_LIBRARIES-NOTFOUND") set(METIS_LIBRARY_DIRS "METIS_LIBRARY_DIRS-NOTFOUND") if(NOT METIS_FIND_QUIETLY) message(STATUS "Looking for metis -- lib metis not found") endif() endif () # check a function to validate the find if(METIS_LIBRARIES) set(REQUIRED_INCDIRS) set(REQUIRED_LIBDIRS) set(REQUIRED_LIBS) # METIS if (METIS_INCLUDE_DIRS) set(REQUIRED_INCDIRS "${METIS_INCLUDE_DIRS}") endif() if (METIS_LIBRARY_DIRS) set(REQUIRED_LIBDIRS "${METIS_LIBRARY_DIRS}") endif() set(REQUIRED_LIBS "${METIS_LIBRARIES}") # m find_library(M_LIBRARY NAMES m) mark_as_advanced(M_LIBRARY) if(M_LIBRARY) list(APPEND REQUIRED_LIBS "-lm") endif() # set required libraries for link set(CMAKE_REQUIRED_INCLUDES "${REQUIRED_INCDIRS}") set(CMAKE_REQUIRED_LIBRARIES) foreach(lib_dir ${REQUIRED_LIBDIRS}) list(APPEND CMAKE_REQUIRED_LIBRARIES "-L${lib_dir}") endforeach() list(APPEND CMAKE_REQUIRED_LIBRARIES "${REQUIRED_LIBS}") string(REGEX REPLACE "^ -" "-" CMAKE_REQUIRED_LIBRARIES "${CMAKE_REQUIRED_LIBRARIES}") # test link unset(METIS_WORKS CACHE) include(CheckFunctionExists) check_function_exists(METIS_NodeND METIS_WORKS) mark_as_advanced(METIS_WORKS) if(NOT METIS_WORKS) if(NOT METIS_FIND_QUIETLY) message(STATUS "Looking for METIS : test of METIS_NodeND with METIS library fails") message(STATUS "CMAKE_REQUIRED_LIBRARIES: ${CMAKE_REQUIRED_LIBRARIES}") message(STATUS "CMAKE_REQUIRED_INCLUDES: ${CMAKE_REQUIRED_INCLUDES}") message(STATUS "Check in CMakeFiles/CMakeError.log to figure out why it fails") endif() endif() set(CMAKE_REQUIRED_INCLUDES) set(CMAKE_REQUIRED_FLAGS) set(CMAKE_REQUIRED_LIBRARIES) endif() if (METIS_LIBRARIES) list(GET METIS_LIBRARIES 0 first_lib) get_filename_component(first_lib_path "${first_lib}" PATH) if (${first_lib_path} MATCHES "/lib(32|64)?$") string(REGEX REPLACE "/lib(32|64)?$" "" not_cached_dir "${first_lib_path}") set(METIS_DIR_FOUND "${not_cached_dir}" CACHE PATH "Installation directory of METIS library" FORCE) else() set(METIS_DIR_FOUND "${first_lib_path}" CACHE PATH "Installation directory of METIS library" FORCE) endif() endif() mark_as_advanced(METIS_DIR) mark_as_advanced(METIS_DIR_FOUND) # check that METIS has been found # --------------------------------- include(FindPackageHandleStandardArgs) find_package_handle_standard_args(METIS DEFAULT_MSG METIS_LIBRARIES METIS_WORKS METIS_INCLUDE_DIRS) # # TODO: Add possibility to check for specific functions in the library # pytorch_sparse-0.6.18/cmake/TorchSparseConfig.cmake.in000066400000000000000000000020411450761466100227340ustar00rootroot00000000000000# 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() pytorch_sparse-0.6.18/conda/000077500000000000000000000000001450761466100157515ustar00rootroot00000000000000pytorch_sparse-0.6.18/conda/pytorch-sparse/000077500000000000000000000000001450761466100207345ustar00rootroot00000000000000pytorch_sparse-0.6.18/conda/pytorch-sparse/README.md000066400000000000000000000001151450761466100222100ustar00rootroot00000000000000``` ./build_conda.sh 3.9 2.1.0 cu118 # python, pytorch and cuda version ``` pytorch_sparse-0.6.18/conda/pytorch-sparse/bld.bat000066400000000000000000000005651450761466100221730ustar00rootroot00000000000000copy "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 pytorch_sparse-0.6.18/conda/pytorch-sparse/build.sh000066400000000000000000000000311450761466100223610ustar00rootroot00000000000000$PYTHON -m pip install . pytorch_sparse-0.6.18/conda/pytorch-sparse/build_conda.sh000077500000000000000000000032051450761466100235360ustar00rootroot00000000000000#!/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 cu121) export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda==12.1.*" ;; cu118) export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda==11.8.*" ;; cu117) export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda==11.7.*" ;; cu116) if [ "${TORCH_VERSION}" = "1.12.0" ]; then export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.6.*" else export CONDA_CUDATOOLKIT_CONSTRAINT="pytorch-cuda==11.6.*" fi ;; 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 [ "${TORCH_VERSION}" = "1.12.0" ] && [ "${CUDA_VERSION}" = "cu116" ]; then conda build . -c pytorch -c pyg -c default -c nvidia -c conda-forge --output-folder "$HOME/conda-bld" else conda build . -c pytorch -c pyg -c default -c nvidia --output-folder "$HOME/conda-bld" fi pytorch_sparse-0.6.18/conda/pytorch-sparse/meta.yaml000066400000000000000000000016571450761466100225570ustar00rootroot00000000000000package: name: pytorch-sparse version: 0.6.18 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 pytorch_sparse-0.6.18/csrc/000077500000000000000000000000001450761466100156175ustar00rootroot00000000000000pytorch_sparse-0.6.18/csrc/convert.cpp000066400000000000000000000020351450761466100200030ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/convert_cpu.h" #ifdef WITH_CUDA #include "cuda/convert_cuda.h" #endif #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__convert_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) { if (ind.device().is_cuda()) { #ifdef WITH_CUDA return ind2ptr_cuda(ind, M); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return ind2ptr_cpu(ind, M); } } SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) { if (ptr.device().is_cuda()) { #ifdef WITH_CUDA return ptr2ind_cuda(ptr, E); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return ptr2ind_cpu(ptr, E); } } static auto registry = torch::RegisterOperators() .op("torch_sparse::ind2ptr", &ind2ptr) .op("torch_sparse::ptr2ind", &ptr2ind); pytorch_sparse-0.6.18/csrc/cpu/000077500000000000000000000000001450761466100164065ustar00rootroot00000000000000pytorch_sparse-0.6.18/csrc/cpu/convert_cpu.cpp000066400000000000000000000027311450761466100214440ustar00rootroot00000000000000#include "convert_cpu.h" #include #include "utils.h" torch::Tensor ind2ptr_cpu(torch::Tensor ind, int64_t M) { CHECK_CPU(ind); auto out = torch::empty(M + 1, ind.options()); auto ind_data = ind.data_ptr(); auto out_data = out.data_ptr(); int64_t numel = ind.numel(); if (numel == 0) return out.zero_(); for (int64_t i = 0; i <= ind_data[0]; i++) out_data[i] = 0; int64_t grain_size = at::internal::GRAIN_SIZE; at::parallel_for(0, numel, grain_size, [&](int64_t begin, int64_t end) { int64_t idx = ind_data[begin], next_idx; for (int64_t i = begin; i < std::min(end, numel - 1); i++) { next_idx = ind_data[i + 1]; for (; idx < next_idx; idx++) out_data[idx + 1] = i + 1; } }); for (int64_t i = ind_data[numel - 1] + 1; i < M + 1; i++) out_data[i] = numel; return out; } torch::Tensor ptr2ind_cpu(torch::Tensor ptr, int64_t E) { CHECK_CPU(ptr); auto out = torch::empty(E, ptr.options()); auto ptr_data = ptr.data_ptr(); auto out_data = out.data_ptr(); int64_t numel = ptr.numel(); int64_t grain_size = at::internal::GRAIN_SIZE; at::parallel_for(0, numel - 1, grain_size, [&](int64_t begin, int64_t end) { int64_t idx = ptr_data[begin], next_idx; for (int64_t i = begin; i < end; i++) { next_idx = ptr_data[i + 1]; for (int64_t e = idx; e < next_idx; e++) out_data[e] = i; idx = next_idx; } }); return out; } pytorch_sparse-0.6.18/csrc/cpu/convert_cpu.h000066400000000000000000000002341450761466100211050ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor ind2ptr_cpu(torch::Tensor ind, int64_t M); torch::Tensor ptr2ind_cpu(torch::Tensor ptr, int64_t E); pytorch_sparse-0.6.18/csrc/cpu/diag_cpu.cpp000066400000000000000000000023051450761466100206650ustar00rootroot00000000000000#include "diag_cpu.h" #include "utils.h" torch::Tensor non_diag_mask_cpu(torch::Tensor row, torch::Tensor col, int64_t M, int64_t N, int64_t k) { CHECK_CPU(row); CHECK_CPU(col); auto E = row.size(0); auto num_diag = k < 0 ? std::min(M + k, N) : std::min(M, N - k); auto row_data = row.data_ptr(); auto col_data = col.data_ptr(); auto mask = torch::zeros({E + num_diag}, row.options().dtype(torch::kBool)); auto mask_data = mask.data_ptr(); int64_t r, c; if (k < 0) { for (int64_t i = 0; i < E; i++) { r = row_data[i], c = col_data[i]; if (r + k < 0) { mask_data[i] = true; } else if (r + k >= N) { mask_data[i + num_diag] = true; } else if (r + k > c) { mask_data[i + r + k] = true; } else if (r + k < c) { mask_data[i + r + k + 1] = true; } } } else { for (int64_t i = 0; i < E; i++) { r = row_data[i], c = col_data[i]; if (r + k >= N) { mask_data[i + num_diag] = true; } else if (r + k > c) { mask_data[i + r] = true; } else if (r + k < c) { mask_data[i + r + 1] = true; } } } return mask; } pytorch_sparse-0.6.18/csrc/cpu/diag_cpu.h000066400000000000000000000002621450761466100203320ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor non_diag_mask_cpu(torch::Tensor row, torch::Tensor col, int64_t M, int64_t N, int64_t k); pytorch_sparse-0.6.18/csrc/cpu/ego_sample_cpu.cpp000066400000000000000000000105371450761466100221020ustar00rootroot00000000000000#include "ego_sample_cpu.h" #include #include "utils.h" #ifdef _WIN32 #include #endif inline torch::Tensor vec2tensor(std::vector vec) { return torch::from_blob(vec.data(), {(int64_t)vec.size()}, at::kLong).clone(); } // Returns `rowptr`, `col`, `n_id`, `e_id`, `ptr`, `root_n_id` std::tuple ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t depth, int64_t num_neighbors, bool replace) { std::vector out_rowptrs(idx.numel() + 1); std::vector out_cols(idx.numel()); std::vector out_n_ids(idx.numel()); std::vector out_e_ids(idx.numel()); auto out_root_n_id = torch::empty({idx.numel()}, at::kLong); out_rowptrs[0] = torch::zeros({1}, at::kLong); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto idx_data = idx.data_ptr(); auto out_root_n_id_data = out_root_n_id.data_ptr(); at::parallel_for(0, idx.numel(), 1, [&](int64_t begin, int64_t end) { int64_t row_start, row_end, row_count, vec_start, vec_end, v, w; for (int64_t g = begin; g < end; g++) { std::set n_id_set; n_id_set.insert(idx_data[g]); std::vector n_ids; n_ids.push_back(idx_data[g]); vec_start = 0, vec_end = n_ids.size(); for (int64_t d = 0; d < depth; d++) { for (int64_t i = vec_start; i < vec_end; i++) { v = n_ids[i]; row_start = rowptr_data[v], row_end = rowptr_data[v + 1]; row_count = row_end - row_start; if (row_count <= num_neighbors) { for (int64_t e = row_start; e < row_end; e++) { w = col_data[e]; n_id_set.insert(w); n_ids.push_back(w); } } else if (replace) { for (int64_t j = 0; j < num_neighbors; j++) { w = col_data[row_start + uniform_randint(row_count)]; n_id_set.insert(w); n_ids.push_back(w); } } else { std::unordered_set perm; for (int64_t j = row_count - num_neighbors; j < row_count; j++) { if (!perm.insert(uniform_randint(j)).second) { perm.insert(j); } } for (int64_t j : perm) { w = col_data[row_start + j]; n_id_set.insert(w); n_ids.push_back(w); } } } vec_start = vec_end; vec_end = n_ids.size(); } n_ids.clear(); std::map n_id_map; std::map::iterator iter; int64_t i = 0; for (int64_t v : n_id_set) { n_ids.push_back(v); n_id_map[v] = i; i++; } out_root_n_id_data[g] = n_id_map[idx_data[g]]; std::vector rowptrs, cols, e_ids; for (int64_t v : n_ids) { row_start = rowptr_data[v], row_end = rowptr_data[v + 1]; for (int64_t e = row_start; e < row_end; e++) { w = col_data[e]; iter = n_id_map.find(w); if (iter != n_id_map.end()) { cols.push_back(iter->second); e_ids.push_back(e); } } rowptrs.push_back(cols.size()); } out_rowptrs[g + 1] = vec2tensor(rowptrs); out_cols[g] = vec2tensor(cols); out_n_ids[g] = vec2tensor(n_ids); out_e_ids[g] = vec2tensor(e_ids); } }); auto out_ptr = torch::empty({idx.numel() + 1}, at::kLong); auto out_ptr_data = out_ptr.data_ptr(); out_ptr_data[0] = 0; int64_t node_cumsum = 0, edge_cumsum = 0; for (int64_t g = 1; g < idx.numel(); g++) { node_cumsum += out_n_ids[g - 1].numel(); edge_cumsum += out_cols[g - 1].numel(); out_rowptrs[g + 1].add_(edge_cumsum); out_cols[g].add_(node_cumsum); out_ptr_data[g] = node_cumsum; out_root_n_id_data[g] += node_cumsum; } node_cumsum += out_n_ids[idx.numel() - 1].numel(); out_ptr_data[idx.numel()] = node_cumsum; return std::make_tuple(torch::cat(out_rowptrs, 0), torch::cat(out_cols, 0), torch::cat(out_n_ids, 0), torch::cat(out_e_ids, 0), out_ptr, out_root_n_id); } pytorch_sparse-0.6.18/csrc/cpu/ego_sample_cpu.h000066400000000000000000000005261450761466100215440ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t depth, int64_t num_neighbors, bool replace); pytorch_sparse-0.6.18/csrc/cpu/hgt_sample_cpu.cpp000066400000000000000000000215701450761466100221110ustar00rootroot00000000000000#include "hgt_sample_cpu.h" #include "utils.h" #ifdef _WIN32 #include #endif #define MAX_NEIGHBORS 50 using namespace std; edge_t split(const rel_t &rel_type) { vector result(3); int start = 0, end; for (int i = 0; i < 3; i++) { end = rel_type.find("__", start); result[i] = rel_type.substr(start, end - start); start = end + 2; } return make_tuple(result[0], result[1], result[2]); } void update_budget_( unordered_map> *budget_dict, const node_t &node_type, const vector &samples, const unordered_map> &to_local_node_dict, const unordered_map &to_edge_type, const c10::Dict &colptr_dict, const c10::Dict &row_dict) { if (samples.empty()) return; for (const auto &kv : colptr_dict) { const auto &rel_type = kv.key(); const auto &edge_type = to_edge_type.at(rel_type); const auto &src_node_type = get<0>(edge_type); const auto &dst_node_type = get<2>(edge_type); if (node_type != dst_node_type) continue; const auto &to_local_src_node = to_local_node_dict.at(src_node_type); const auto *colptr_data = kv.value().data_ptr(); const auto *row_data = row_dict.at(rel_type).data_ptr(); auto &src_budget = budget_dict->at(src_node_type); for (const auto &w : samples) { const auto &col_start = colptr_data[w], &col_end = colptr_data[w + 1]; if (col_end - col_start > MAX_NEIGHBORS) { // There might be same neighbors with large neighborhood sizes. // In order to prevent that we fill our budget with many values of low // probability, we instead sample a fixed amount without replacement: auto indices = choice(col_end - col_start, MAX_NEIGHBORS, false); auto *indices_data = indices.data_ptr(); for (int64_t i = 0; i < indices.numel(); i++) { const auto &v = row_data[col_start + indices_data[i]]; // Only add the neighbor in case we have not yet seen it before: if (to_local_src_node.find(v) == to_local_src_node.end()) src_budget[v] += 1.f / float(MAX_NEIGHBORS); } } else if (col_end != col_start) { const auto inv_deg = 1.f / float(col_end - col_start); for (int64_t i = col_start; i < col_end; i++) { const auto &v = row_data[i]; // Only add the neighbor in case we have not yet seen it before: if (to_local_src_node.find(v) == to_local_src_node.end()) src_budget[v] += inv_deg; } } } } } vector sample_from(const unordered_map &budget, const int64_t num_samples) { vector indices; vector weights; indices.reserve(budget.size()); weights.reserve(budget.size()); for (const auto &kv : budget) { indices.push_back(kv.first); weights.push_back(kv.second * kv.second); } const auto weight = from_vector(weights, true); const auto sample = choice(budget.size(), num_samples, false, weight); const auto *sample_data = sample.data_ptr(); vector out(sample.numel()); for (int64_t i = 0; i < sample.numel(); i++) { out[i] = indices[sample_data[i]]; } return out; } tuple, c10::Dict, c10::Dict, c10::Dict> hgt_sample_cpu(const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_samples_dict, const int64_t num_hops) { // Create a mapping to convert single string relations to edge type triplets: unordered_map to_edge_type; for (const auto &kv : colptr_dict) { const auto &rel_type = kv.key(); to_edge_type[rel_type] = split(rel_type); } // Initialize some necessary data structures for the sampling process: unordered_map> nodes_dict; unordered_map> to_local_node_dict; unordered_map> budget_dict; for (const auto &kv : num_samples_dict) { const auto &node_type = kv.key(); nodes_dict[node_type]; to_local_node_dict[node_type]; budget_dict[node_type]; } // Add the input nodes to the sampled output nodes (line 1): for (const auto &kv : input_node_dict) { const auto &node_type = kv.key(); const auto &input_node = kv.value(); const auto *input_node_data = input_node.data_ptr(); auto &nodes = nodes_dict.at(node_type); auto &to_local_node = to_local_node_dict.at(node_type); for (int64_t i = 0; i < input_node.numel(); i++) { const auto &v = input_node_data[i]; nodes.push_back(v); to_local_node[v] = i; } } // Update the budget based on the initial input set (line 3-5): for (const auto &kv : nodes_dict) { const auto &node_type = kv.first; const auto &last_samples = kv.second; update_budget_(&budget_dict, node_type, last_samples, to_local_node_dict, to_edge_type, colptr_dict, row_dict); } for (int64_t ell = 0; ell < num_hops; ell++) { unordered_map> samples_dict; for (auto &kv : budget_dict) { const auto &node_type = kv.first; auto &budget = kv.second; const auto num_samples = num_samples_dict.at(node_type)[ell]; // Sample `num_samples` nodes, according to the budget (line 9-11): const auto samples = sample_from(budget, num_samples); samples_dict[node_type] = samples; // Add samples to the sampled output nodes, and erase them from the budget // (line 13/15): auto &nodes = nodes_dict.at(node_type); auto &to_local_node = to_local_node_dict.at(node_type); for (const auto &v : samples) { to_local_node[v] = nodes.size(); nodes.push_back(v); budget.erase(v); } } if (ell < num_hops - 1) { // Add neighbors of newly sampled nodes to the budget (line 14): // Note that we do not need to update the budget in the last iteration. for (const auto &kv : samples_dict) { const auto &node_type = kv.first; const auto &last_samples = kv.second; update_budget_(&budget_dict, node_type, last_samples, to_local_node_dict, to_edge_type, colptr_dict, row_dict); } } } c10::Dict out_node_dict; c10::Dict out_row_dict; c10::Dict out_col_dict; c10::Dict out_edge_dict; // Reconstruct the sampled adjacency matrix among the sampled nodes (line 19): for (const auto &kv : colptr_dict) { const auto &rel_type = kv.key(); const auto &edge_type = to_edge_type.at(rel_type); const auto &src_node_type = get<0>(edge_type); const auto &dst_node_type = get<2>(edge_type); const auto *colptr_data = kv.value().data_ptr(); const auto *row_data = row_dict.at(rel_type).data_ptr(); const auto &dst_nodes = nodes_dict.at(dst_node_type); const auto &to_local_src_node = to_local_node_dict.at(src_node_type); vector rows, cols, edges; for (int64_t i = 0; i < (int64_t)dst_nodes.size(); i++) { const auto &w = dst_nodes[i]; const auto &col_start = colptr_data[w], &col_end = colptr_data[w + 1]; if (col_end - col_start > MAX_NEIGHBORS) { auto indices = choice(col_end - col_start, MAX_NEIGHBORS, false); auto *indices_data = indices.data_ptr(); for (int64_t j = 0; j < indices.numel(); j++) { const auto &v = row_data[col_start + indices_data[j]]; if (to_local_src_node.find(v) != to_local_src_node.end()) { rows.push_back(to_local_src_node.at(v)); cols.push_back(i); edges.push_back(col_start + j); } } } else { for (int64_t j = col_start; j < col_end; j++) { const auto &v = row_data[j]; if (to_local_src_node.find(v) != to_local_src_node.end()) { rows.push_back(to_local_src_node.at(v)); cols.push_back(i); edges.push_back(j); } } } } out_row_dict.insert(rel_type, from_vector(rows)); out_col_dict.insert(rel_type, from_vector(cols)); out_edge_dict.insert(rel_type, from_vector(edges)); } // Generate tensor-valued output node dictionary (line 20): for (const auto &kv : nodes_dict) { const auto &node_type = kv.first; const auto &nodes = kv.second; if (!nodes.empty()) out_node_dict.insert(node_type, from_vector(nodes)); } return make_tuple(out_node_dict, out_row_dict, out_col_dict, out_edge_dict); } pytorch_sparse-0.6.18/csrc/cpu/hgt_sample_cpu.h000066400000000000000000000012021450761466100215440ustar00rootroot00000000000000#pragma once #include "../extensions.h" typedef std::string node_t; typedef std::string rel_t; typedef std::tuple edge_t; std::tuple, c10::Dict, c10::Dict, c10::Dict> hgt_sample_cpu(const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_samples_dict, const int64_t num_hops); pytorch_sparse-0.6.18/csrc/cpu/metis_cpu.cpp000066400000000000000000000075451450761466100211150ustar00rootroot00000000000000#include "metis_cpu.h" #ifdef WITH_METIS #include #endif #ifdef WITH_MTMETIS #include #endif #include "utils.h" torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive) { #ifdef WITH_METIS CHECK_CPU(rowptr); CHECK_CPU(col); if (optional_value.has_value()) { CHECK_CPU(optional_value.value()); CHECK_INPUT(optional_value.value().dim() == 1); CHECK_INPUT(optional_value.value().numel() == col.numel()); } if (optional_node_weight.has_value()) { CHECK_CPU(optional_node_weight.value()); CHECK_INPUT(optional_node_weight.value().dim() == 1); CHECK_INPUT(optional_node_weight.value().numel() == rowptr.numel() - 1); } int64_t nvtxs = rowptr.numel() - 1; int64_t ncon = 1; auto *xadj = rowptr.data_ptr(); auto *adjncy = col.data_ptr(); int64_t *adjwgt = NULL; if (optional_value.has_value()) adjwgt = optional_value.value().data_ptr(); int64_t *vwgt = NULL; if (optional_node_weight.has_value()) vwgt = optional_node_weight.value().data_ptr(); int64_t objval = -1; auto part = torch::empty({nvtxs}, rowptr.options()); auto part_data = part.data_ptr(); if (recursive) { METIS_PartGraphRecursive(&nvtxs, &ncon, xadj, adjncy, vwgt, NULL, adjwgt, &num_parts, NULL, NULL, NULL, &objval, part_data); } else { METIS_PartGraphKway(&nvtxs, &ncon, xadj, adjncy, vwgt, NULL, adjwgt, &num_parts, NULL, NULL, NULL, &objval, part_data); } return part; #else AT_ERROR("Not compiled with METIS support"); #endif } // needs mt-metis installed via: // ./configure --shared --edges64bit --vertices64bit --weights64bit // --partitions64bit torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive, int64_t num_workers) { #ifdef WITH_MTMETIS CHECK_CPU(rowptr); CHECK_CPU(col); if (optional_value.has_value()) { CHECK_CPU(optional_value.value()); CHECK_INPUT(optional_value.value().dim() == 1); CHECK_INPUT(optional_value.value().numel() == col.numel()); } if (optional_node_weight.has_value()) { CHECK_CPU(optional_node_weight.value()); CHECK_INPUT(optional_node_weight.value().dim() == 1); CHECK_INPUT(optional_node_weight.value().numel() == rowptr.numel() - 1); } mtmetis_vtx_type nvtxs = rowptr.numel() - 1; mtmetis_vtx_type ncon = 1; mtmetis_adj_type *xadj = (mtmetis_adj_type *)rowptr.data_ptr(); mtmetis_vtx_type *adjncy = (mtmetis_vtx_type *)col.data_ptr(); mtmetis_wgt_type *adjwgt = NULL; if (optional_value.has_value()) adjwgt = optional_value.value().data_ptr(); mtmetis_wgt_type *vwgt = NULL; if (optional_node_weight.has_value()) vwgt = optional_node_weight.value().data_ptr(); mtmetis_pid_type nparts = num_parts; mtmetis_wgt_type objval = -1; auto part = torch::empty({nvtxs}, rowptr.options()); mtmetis_pid_type *part_data = (mtmetis_pid_type *)part.data_ptr(); double *opts = mtmetis_init_options(); opts[MTMETIS_OPTION_NTHREADS] = num_workers; if (recursive) { MTMETIS_PartGraphRecursive(&nvtxs, &ncon, xadj, adjncy, vwgt, NULL, adjwgt, &nparts, NULL, NULL, opts, &objval, part_data); } else { MTMETIS_PartGraphKway(&nvtxs, &ncon, xadj, adjncy, vwgt, NULL, adjwgt, &nparts, NULL, NULL, opts, &objval, part_data); } return part; #else AT_ERROR("Not compiled with MTMETIS support"); #endif } pytorch_sparse-0.6.18/csrc/cpu/metis_cpu.h000066400000000000000000000011441450761466100205470ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive); torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive, int64_t num_workers); pytorch_sparse-0.6.18/csrc/cpu/neighbor_sample_cpu.cpp000066400000000000000000000472671450761466100231370ustar00rootroot00000000000000#include "neighbor_sample_cpu.h" #include "utils.h" #ifdef _WIN32 #include #endif using namespace std; namespace { typedef phmap::flat_hash_map, int64_t> temporarl_edge_dict; template tuple sample(const torch::Tensor &colptr, const torch::Tensor &row, const torch::Tensor &input_node, const vector num_neighbors) { // Initialize some data structures for the sampling process: vector samples; phmap::flat_hash_map to_local_node; auto *colptr_data = colptr.data_ptr(); auto *row_data = row.data_ptr(); auto *input_node_data = input_node.data_ptr(); for (int64_t i = 0; i < input_node.numel(); i++) { const auto &v = input_node_data[i]; samples.push_back(v); to_local_node.insert({v, i}); } vector rows, cols, edges; int64_t begin = 0, end = samples.size(); for (int64_t ell = 0; ell < (int64_t)num_neighbors.size(); ell++) { const auto &num_samples = num_neighbors[ell]; for (int64_t i = begin; i < end; i++) { const auto &w = samples[i]; const auto &col_start = colptr_data[w]; const auto &col_end = colptr_data[w + 1]; const auto col_count = col_end - col_start; if (col_count == 0) continue; if ((num_samples < 0) || (!replace && (num_samples >= col_count))) { for (int64_t offset = col_start; offset < col_end; offset++) { const int64_t &v = row_data[offset]; const auto res = to_local_node.insert({v, samples.size()}); if (res.second) samples.push_back(v); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } } else if (replace) { for (int64_t j = 0; j < num_samples; j++) { 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) samples.push_back(v); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } } else { unordered_set rnd_indices; for (int64_t j = col_count - num_samples; j < col_count; 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]; const auto res = to_local_node.insert({v, samples.size()}); if (res.second) samples.push_back(v); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } } } begin = end, end = samples.size(); } if (!directed) { phmap::flat_hash_map::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]; const auto &col_end = colptr_data[w + 1]; for (int64_t offset = col_start; offset < col_end; offset++) { const auto &v = row_data[offset]; iter = to_local_node.find(v); if (iter != to_local_node.end()) { rows.push_back(iter->second); cols.push_back(i); edges.push_back(offset); } } } } return make_tuple(from_vector(samples), from_vector(rows), from_vector(cols), from_vector(edges)); } inline bool satisfy_time(const c10::Dict &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()[src_node] <= dst_time; } catch (const std::out_of_range& e) { // If no time is given, fall back to normal sampling return true; } } template tuple, c10::Dict, c10::Dict, c10::Dict> hetero_sample(const vector &node_types, const vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const c10::Dict &node_time_dict, const int64_t num_hops) { // Create a mapping to convert single string relations to edge type triplets: phmap::flat_hash_map 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: phmap::flat_hash_map> samples_dict; phmap::flat_hash_map>> temp_samples_dict; phmap::flat_hash_map> to_local_node_dict; phmap::flat_hash_map temp_to_local_node_dict; phmap::flat_hash_map> 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]; } phmap::flat_hash_map> rows_dict, cols_dict, edges_dict; for (const auto &kv : colptr_dict) { const auto &rel_type = kv.key(); rows_dict[rel_type]; cols_dict[rel_type]; edges_dict[rel_type]; } // Add the input nodes to the output nodes: for (const auto &kv : input_node_dict) { const auto &node_type = kv.key(); const torch::Tensor &input_node = kv.value(); const auto *input_node_data = input_node.data_ptr(); 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(); } 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]); } } phmap::flat_hash_map> 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()}; } vector all_rel_types; for (const auto &kv : num_neighbors_dict) { 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 = 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 torch::Tensor &colptr = colptr_dict.at(rel_type); const auto *colptr_data = colptr.data_ptr(); const torch::Tensor &row = row_dict.at(rel_type); const auto *row_data = row.data_ptr(); 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 = 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; if (col_count == 0) 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); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } } } else if (replace) { // 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); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } num_neighbors += 1; } } else { // Sample without replacement: unordered_set rnd_indices; for (int64_t j = col_count - num_samples; j < col_count; 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); if (directed) { cols.push_back(i); rows.push_back(res.first->second); edges.push_back(offset); } } } } } } if (temporal) { for (const auto &kv : temp_samples_dict) { slice_dict[kv.first] = {slice_dict.at(kv.first).second, kv.second.size()}; } } else { for (const auto &kv : samples_dict) slice_dict[kv.first] = {slice_dict.at(kv.first).second, kv.second.size()}; } } // Temporal sample disable undirected assert(!(temporal && !directed)); if (!directed) { // Construct the subgraph among the sampled nodes: phmap::flat_hash_map::iterator iter; for (const auto &kv : colptr_dict) { const auto &rel_type = kv.key(); 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 &dst_samples = samples_dict.at(dst_node_type); auto &to_local_src_node = to_local_node_dict.at(src_node_type); const auto *colptr_data = ((torch::Tensor)kv.value()).data_ptr(); const auto *row_data = ((torch::Tensor)row_dict.at(rel_type)).data_ptr(); auto &rows = rows_dict.at(rel_type); auto &cols = cols_dict.at(rel_type); auto &edges = edges_dict.at(rel_type); for (int64_t i = 0; i < (int64_t)dst_samples.size(); i++) { const auto &w = dst_samples[i]; const auto &col_start = colptr_data[w]; const auto &col_end = colptr_data[w + 1]; for (int64_t offset = col_start; offset < col_end; offset++) { const auto &v = row_data[offset]; iter = to_local_src_node.find(v); if (iter != to_local_src_node.end()) { rows.push_back(iter->second); cols.push_back(i); edges.push_back(offset); } } } } } // 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(samples_dict), from_vector(rows_dict), from_vector(cols_dict), from_vector(edges_dict)); } } // namespace tuple neighbor_sample_cpu(const torch::Tensor &colptr, const torch::Tensor &row, const torch::Tensor &input_node, const vector num_neighbors, const bool replace, const bool directed) { if (replace && directed) { return sample(colptr, row, input_node, num_neighbors); } else if (replace && !directed) { return sample(colptr, row, input_node, num_neighbors); } else if (!replace && directed) { return sample(colptr, row, input_node, num_neighbors); } else { return sample(colptr, row, input_node, num_neighbors); } } tuple, c10::Dict, c10::Dict, c10::Dict> hetero_neighbor_sample_cpu( const vector &node_types, const vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const int64_t num_hops, const bool replace, const bool directed) { c10::Dict node_time_dict; // Empty dictionary. if (replace && directed) { return hetero_sample( 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( 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( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, node_time_dict, num_hops); } else { return hetero_sample( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, node_time_dict, num_hops); } } tuple, c10::Dict, c10::Dict, c10::Dict> hetero_temporal_neighbor_sample_cpu( const vector &node_types, const vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const c10::Dict &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( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, node_time_dict, num_hops); } else { return hetero_sample( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, node_time_dict, num_hops); } } pytorch_sparse-0.6.18/csrc/cpu/neighbor_sample_cpu.h000066400000000000000000000032641450761466100225710ustar00rootroot00000000000000#pragma once #include "../extensions.h" typedef std::string node_t; typedef std::tuple edge_t; typedef std::string rel_t; std::tuple neighbor_sample_cpu(const torch::Tensor &colptr, const torch::Tensor &row, const torch::Tensor &input_node, const std::vector num_neighbors, const bool replace, const bool directed); std::tuple, c10::Dict, c10::Dict, c10::Dict> hetero_neighbor_sample_cpu( const std::vector &node_types, const std::vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const int64_t num_hops, const bool replace, const bool directed); std::tuple, c10::Dict, c10::Dict, c10::Dict> hetero_temporal_neighbor_sample_cpu( const std::vector &node_types, const std::vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const c10::Dict &node_time_dict, const int64_t num_hops, const bool replace, const bool directed); pytorch_sparse-0.6.18/csrc/cpu/reducer.h000066400000000000000000000074121450761466100202140ustar00rootroot00000000000000#pragma once #include #include enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX }; const std::map reduce2REDUCE = { {"sum", SUM}, {"mean", MEAN}, {"mul", MUL}, {"div", DIV}, {"min", MIN}, {"max", MAX}, }; #define AT_DISPATCH_REDUCTION_TYPES(reduce, ...) \ [&] { \ switch (reduce2REDUCE.at(reduce)) { \ case SUM: { \ static constexpr ReductionType REDUCE = SUM; \ return __VA_ARGS__(); \ } \ case MEAN: { \ static constexpr ReductionType REDUCE = MEAN; \ return __VA_ARGS__(); \ } \ case MUL: { \ static constexpr ReductionType REDUCE = MUL; \ return __VA_ARGS__(); \ } \ case DIV: { \ static constexpr ReductionType REDUCE = DIV; \ return __VA_ARGS__(); \ } \ case MIN: { \ static constexpr ReductionType REDUCE = MIN; \ return __VA_ARGS__(); \ } \ case MAX: { \ static constexpr ReductionType REDUCE = MAX; \ return __VA_ARGS__(); \ } \ } \ }() template struct Reducer { static inline scalar_t init() { if (REDUCE == MUL || REDUCE == DIV) return (scalar_t)1; else if (REDUCE == MIN) return std::numeric_limits::max(); else if (REDUCE == MAX) return std::numeric_limits::lowest(); else return (scalar_t)0; } static inline void update(scalar_t *val, scalar_t new_val, int64_t *arg, int64_t new_arg) { if (REDUCE == SUM || REDUCE == MEAN) *val = *val + new_val; else if (REDUCE == MUL) *val = *val * new_val; else if (REDUCE == DIV) *val = *val / new_val; else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) { *val = new_val; *arg = new_arg; } } static inline void write(scalar_t *address, scalar_t val, int64_t *arg_address, int64_t arg, int count) { if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV) *address = val; else if (REDUCE == MEAN) *address = val / (scalar_t)(count > 0 ? count : 1); else if (REDUCE == MIN || REDUCE == MAX) { if (count > 0) { *address = val; *arg_address = arg; } else *address = (scalar_t)0; } } }; pytorch_sparse-0.6.18/csrc/cpu/relabel_cpu.cpp000066400000000000000000000102531450761466100213700ustar00rootroot00000000000000#include "relabel_cpu.h" #include "utils.h" std::tuple relabel_cpu(torch::Tensor col, torch::Tensor idx) { CHECK_CPU(col); CHECK_CPU(idx); CHECK_INPUT(idx.dim() == 1); auto col_data = col.data_ptr(); auto idx_data = idx.data_ptr(); std::vector cols; std::vector n_ids; std::unordered_map n_id_map; int64_t i; for (int64_t n = 0; n < idx.size(0); n++) { i = idx_data[n]; n_id_map[i] = n; n_ids.push_back(i); } int64_t c; for (int64_t e = 0; e < col.size(0); e++) { c = col_data[e]; if (n_id_map.count(c) == 0) { n_id_map[c] = n_ids.size(); n_ids.push_back(c); } cols.push_back(n_id_map[c]); } int64_t n_len = n_ids.size(), e_len = cols.size(); auto out_col = torch::from_blob(cols.data(), {e_len}, col.options()).clone(); auto out_idx = torch::from_blob(n_ids.data(), {n_len}, col.options()).clone(); return std::make_tuple(out_col, out_idx); } std::tuple, torch::Tensor> relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor idx, bool bipartite) { CHECK_CPU(rowptr); CHECK_CPU(col); if (optional_value.has_value()) { CHECK_CPU(optional_value.value()); CHECK_INPUT(optional_value.value().dim() == 1); } CHECK_CPU(idx); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto idx_data = idx.data_ptr(); std::vector n_ids; std::unordered_map n_id_map; std::unordered_map::iterator it; auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options()); auto out_rowptr_data = out_rowptr.data_ptr(); out_rowptr_data[0] = 0; int64_t v, w, c, row_start, row_end, offset = 0; for (int64_t i = 0; i < idx.numel(); i++) { v = idx_data[i]; n_id_map[v] = i; offset += rowptr_data[v + 1] - rowptr_data[v]; out_rowptr_data[i + 1] = offset; } auto out_col = torch::empty({offset}, col.options()); auto out_col_data = out_col.data_ptr(); torch::optional out_value = torch::nullopt; if (optional_value.has_value()) { 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(); auto out_value_data = out_value.value().data_ptr(); offset = 0; for (int64_t i = 0; i < idx.numel(); i++) { v = idx_data[i]; row_start = rowptr_data[v], row_end = rowptr_data[v + 1]; for (int64_t j = row_start; j < row_end; j++) { w = col_data[j]; it = n_id_map.find(w); if (it == n_id_map.end()) { c = idx.numel() + n_ids.size(); n_id_map[w] = c; n_ids.push_back(w); out_col_data[offset] = c; } else { out_col_data[offset] = it->second; } out_value_data[offset] = value_data[j]; offset++; } } }); } else { offset = 0; for (int64_t i = 0; i < idx.numel(); i++) { v = idx_data[i]; row_start = rowptr_data[v], row_end = rowptr_data[v + 1]; for (int64_t j = row_start; j < row_end; j++) { w = col_data[j]; it = n_id_map.find(w); if (it == n_id_map.end()) { c = idx.numel() + n_ids.size(); n_id_map[w] = c; n_ids.push_back(w); out_col_data[offset] = c; } else { out_col_data[offset] = it->second; } offset++; } } } if (!bipartite) out_rowptr = torch::cat( {out_rowptr, torch::full({(int64_t)n_ids.size()}, out_col.numel(), rowptr.options())}); idx = torch::cat({idx, torch::from_blob(n_ids.data(), {(int64_t)n_ids.size()}, idx.options())}); return std::make_tuple(out_rowptr, out_col, out_value, idx); } pytorch_sparse-0.6.18/csrc/cpu/relabel_cpu.h000066400000000000000000000007271450761466100210420ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple relabel_cpu(torch::Tensor col, torch::Tensor idx); std::tuple, torch::Tensor> relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor idx, bool bipartite); pytorch_sparse-0.6.18/csrc/cpu/rw_cpu.cpp000066400000000000000000000023401450761466100204100ustar00rootroot00000000000000#include "rw_cpu.h" #include "utils.h" torch::Tensor random_walk_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length) { CHECK_CPU(rowptr); CHECK_CPU(col); CHECK_CPU(start); CHECK_INPUT(rowptr.dim() == 1); CHECK_INPUT(col.dim() == 1); CHECK_INPUT(start.dim() == 1); auto rand = torch::rand({start.size(0), walk_length}, start.options().dtype(torch::kFloat)); auto L = walk_length + 1; auto out = torch::full({start.size(0), L}, -1, start.options()); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto start_data = start.data_ptr(); auto rand_data = rand.data_ptr(); auto out_data = out.data_ptr(); for (auto n = 0; n < start.size(0); n++) { auto cur = start_data[n]; out_data[n * L] = cur; int64_t row_start, row_end; for (auto l = 0; l < walk_length; l++) { row_start = rowptr_data[cur]; row_end = rowptr_data[cur + 1]; cur = col_data[row_start + int64_t(rand_data[n * walk_length + l] * (row_end - row_start))]; out_data[n * L + l + 1] = cur; } } return out; } pytorch_sparse-0.6.18/csrc/cpu/rw_cpu.h000066400000000000000000000002721450761466100200570ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor random_walk_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length); pytorch_sparse-0.6.18/csrc/cpu/saint_cpu.cpp000066400000000000000000000026631450761466100211060ustar00rootroot00000000000000#include "saint_cpu.h" #include "utils.h" std::tuple subgraph_cpu(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row, torch::Tensor col) { CHECK_CPU(idx); CHECK_CPU(rowptr); CHECK_CPU(col); CHECK_INPUT(idx.dim() == 1); CHECK_INPUT(rowptr.dim() == 1); CHECK_INPUT(col.dim() == 1); auto assoc = torch::full({rowptr.size(0) - 1}, -1, idx.options()); assoc.index_copy_(0, idx, torch::arange(idx.size(0), idx.options())); auto idx_data = idx.data_ptr(); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto assoc_data = assoc.data_ptr(); std::vector rows, cols, indices; int64_t v, w, w_new, row_start, row_end; for (int64_t v_new = 0; v_new < idx.size(0); v_new++) { v = idx_data[v_new]; row_start = rowptr_data[v]; row_end = rowptr_data[v + 1]; for (int64_t j = row_start; j < row_end; j++) { w = col_data[j]; w_new = assoc_data[w]; if (w_new > -1) { rows.push_back(v_new); cols.push_back(w_new); indices.push_back(j); } } } int64_t length = rows.size(); row = torch::from_blob(rows.data(), {length}, row.options()).clone(); col = torch::from_blob(cols.data(), {length}, row.options()).clone(); idx = torch::from_blob(indices.data(), {length}, row.options()).clone(); return std::make_tuple(row, col, idx); } pytorch_sparse-0.6.18/csrc/cpu/saint_cpu.h000066400000000000000000000003141450761466100205420ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple subgraph_cpu(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row, torch::Tensor col); pytorch_sparse-0.6.18/csrc/cpu/sample_cpu.cpp000066400000000000000000000105121450761466100212410ustar00rootroot00000000000000#include "sample_cpu.h" #include "utils.h" #ifdef _WIN32 #include #endif // Returns `rowptr`, `col`, `n_id`, `e_id` std::tuple sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t num_neighbors, bool replace) { CHECK_CPU(rowptr); CHECK_CPU(col); CHECK_CPU(idx); CHECK_INPUT(idx.dim() == 1); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto idx_data = idx.data_ptr(); auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options()); auto out_rowptr_data = out_rowptr.data_ptr(); out_rowptr_data[0] = 0; std::vector>> cols; // col, e_id std::vector n_ids; std::unordered_map n_id_map; int64_t i; for (int64_t n = 0; n < idx.numel(); n++) { i = idx_data[n]; cols.push_back(std::vector>()); n_id_map[i] = n; n_ids.push_back(i); } int64_t n, c, e, row_start, row_end, row_count; if (num_neighbors < 0) { // No sampling ====================================== for (int64_t i = 0; i < idx.numel(); i++) { n = idx_data[i]; row_start = rowptr_data[n], row_end = rowptr_data[n + 1]; row_count = row_end - row_start; for (int64_t j = 0; j < row_count; j++) { e = row_start + j; c = col_data[e]; if (n_id_map.count(c) == 0) { n_id_map[c] = n_ids.size(); n_ids.push_back(c); } cols[i].push_back(std::make_tuple(n_id_map[c], e)); } out_rowptr_data[i + 1] = out_rowptr_data[i] + cols[i].size(); } } else if (replace) { // Sample with replacement =============================== for (int64_t i = 0; i < idx.numel(); i++) { n = idx_data[i]; row_start = rowptr_data[n], row_end = rowptr_data[n + 1]; row_count = row_end - row_start; if (row_count > 0) { for (int64_t j = 0; j < num_neighbors; j++) { e = row_start + uniform_randint(row_count); c = col_data[e]; if (n_id_map.count(c) == 0) { n_id_map[c] = n_ids.size(); n_ids.push_back(c); } cols[i].push_back(std::make_tuple(n_id_map[c], e)); } } out_rowptr_data[i + 1] = out_rowptr_data[i] + cols[i].size(); } } else { // Sample without replacement via Robert Floyd algorithm ============ for (int64_t i = 0; i < idx.numel(); i++) { n = idx_data[i]; row_start = rowptr_data[n], row_end = rowptr_data[n + 1]; row_count = row_end - row_start; std::unordered_set perm; if (row_count <= num_neighbors) { for (int64_t j = 0; j < row_count; j++) perm.insert(j); } 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(uniform_randint(j)).second) perm.insert(j); } } for (const int64_t &p : perm) { e = row_start + p; c = col_data[e]; if (n_id_map.count(c) == 0) { n_id_map[c] = n_ids.size(); n_ids.push_back(c); } cols[i].push_back(std::make_tuple(n_id_map[c], e)); } out_rowptr_data[i + 1] = out_rowptr_data[i] + cols[i].size(); } } int64_t N = n_ids.size(); 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_data = out_col.data_ptr(); auto out_e_id = torch::empty({E}, col.options()); auto out_e_id_data = out_e_id.data_ptr(); i = 0; for (std::vector> &col_vec : cols) { std::sort(col_vec.begin(), col_vec.end(), [](const std::tuple &a, const std::tuple &b) -> bool { return std::get<0>(a) < std::get<0>(b); }); for (const std::tuple &value : col_vec) { out_col_data[i] = std::get<0>(value); out_e_id_data[i] = std::get<1>(value); i += 1; } } return std::make_tuple(out_rowptr, out_col, out_n_id, out_e_id); } pytorch_sparse-0.6.18/csrc/cpu/sample_cpu.h000066400000000000000000000003611450761466100207070ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t num_neighbors, bool replace); pytorch_sparse-0.6.18/csrc/cpu/spmm_cpu.cpp000066400000000000000000000116541450761466100207440ustar00rootroot00000000000000#include "spmm_cpu.h" #include #include "reducer.h" #include "utils.h" std::tuple> spmm_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor mat, std::string reduce) { CHECK_CPU(rowptr); CHECK_CPU(col); if (optional_value.has_value()) CHECK_CPU(optional_value.value()); CHECK_CPU(mat); CHECK_INPUT(rowptr.dim() == 1); CHECK_INPUT(col.dim() == 1); if (optional_value.has_value()) { CHECK_INPUT(optional_value.value().dim() == 1); CHECK_INPUT(optional_value.value().size(0) == col.size(0)); } CHECK_INPUT(mat.dim() >= 2); mat = mat.contiguous(); auto sizes = mat.sizes().vec(); sizes[mat.dim() - 2] = rowptr.numel() - 1; auto out = torch::empty(sizes, mat.options()); torch::optional arg_out = torch::nullopt; int64_t *arg_out_data = nullptr; if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) { arg_out = torch::full_like(out, col.numel(), rowptr.options()); arg_out_data = arg_out.value().data_ptr(); } auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto M = rowptr.numel() - 1; auto N = mat.size(-2); auto K = mat.size(-1); auto B = mat.numel() / (N * K); AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, mat.scalar_type(), "spmm_cpu", [&] { scalar_t *value_data = nullptr; auto mat_data = mat.data_ptr(); auto out_data = out.data_ptr(); AT_DISPATCH_REDUCTION_TYPES(reduce, [&] { AT_DISPATCH_HAS_VALUE(optional_value, [&] { if (HAS_VALUE) { value_data = optional_value.value().data_ptr(); } int64_t grain_size = at::internal::GRAIN_SIZE / (K * std::max(col.numel() / M, (int64_t)1)); at::parallel_for(0, B * M, grain_size, [&](int64_t begin, int64_t end) { scalar_t val; std::vector vals(K); int64_t row_start, row_end, b, m, c; std::vector args(K); for (auto i = begin; i < end; i++) { b = i / M, m = i % M; row_start = rowptr_data[m], row_end = rowptr_data[m + 1]; for (auto k = 0; k < K; k++) vals[k] = Reducer::init(); auto offset = b * N * K; for (auto e = row_start; e < row_end; e++) { c = col_data[e]; if (HAS_VALUE) val = value_data[e]; for (auto k = 0; k < K; k++) { if (HAS_VALUE) Reducer::update( &vals[k], val * mat_data[offset + c * K + k], &args[k], e); else Reducer::update( &vals[k], mat_data[offset + c * K + k], &args[k], e); } } offset = b * M * K + m * K; for (auto k = 0; k < K; k++) Reducer::write(out_data + offset + k, vals[k], arg_out_data + offset + k, args[k], row_end - row_start); } }); }); }); }); return std::make_tuple(out, arg_out); } torch::Tensor spmm_value_bw_cpu(torch::Tensor row, torch::Tensor rowptr, torch::Tensor col, torch::Tensor mat, torch::Tensor grad, std::string reduce) { CHECK_CPU(row); CHECK_CPU(rowptr); CHECK_CPU(col); CHECK_CPU(mat); CHECK_CPU(grad); mat = mat.contiguous(); grad = grad.contiguous(); auto M = grad.size(-2); auto N = mat.size(-2); auto E = row.numel(); auto K = mat.size(-1); auto B = mat.numel() / (N * K); auto out = torch::zeros({row.numel()}, grad.options()); auto row_data = row.data_ptr(); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, mat.scalar_type(), "spmm_value_bw_cpu", [&] { auto mat_data = mat.data_ptr(); auto grad_data = grad.data_ptr(); auto out_data = out.data_ptr(); scalar_t val; int64_t row, col; AT_DISPATCH_REDUCTION_TYPES(reduce, [&] { for (int b = 0; b < B; b++) { for (int e = 0; e < E; e++) { row = row_data[e], col = col_data[e], val = (scalar_t)0; for (int k = 0; k < K; k++) { val += mat_data[b * N * K + col * K + k] * grad_data[b * M * K + row * K + k]; } if (REDUCE == MEAN) { int row_start = rowptr_data[row], row_end = rowptr_data[row + 1]; val /= (scalar_t)std::max(row_end - row_start, 1); } out_data[e] += val; } } }); }); return out; } pytorch_sparse-0.6.18/csrc/cpu/spmm_cpu.h000066400000000000000000000007301450761466100204020ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple> spmm_cpu(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor mat, std::string reduce); torch::Tensor spmm_value_bw_cpu(torch::Tensor row, torch::Tensor rowptr, torch::Tensor col, torch::Tensor mat, torch::Tensor grad, std::string reduce); pytorch_sparse-0.6.18/csrc/cpu/utils.h000066400000000000000000000107361450761466100177260ustar00rootroot00000000000000#pragma once #include "../extensions.h" #include "parallel_hashmap/phmap.h" #define CHECK_CPU(x) AT_ASSERTM(x.device().is_cpu(), #x " must be CPU tensor") #define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch") #define CHECK_LT(low, high) AT_ASSERTM(low < high, "low must be smaller than high") #define AT_DISPATCH_HAS_VALUE(optional_value, ...) \ [&] { \ if (optional_value.has_value()) { \ const bool HAS_VALUE = true; \ return __VA_ARGS__(); \ } else { \ const bool HAS_VALUE = false; \ return __VA_ARGS__(); \ } \ }() template inline torch::Tensor from_vector(const std::vector &vec, bool inplace = false) { const auto size = (int64_t)vec.size(); const auto out = torch::from_blob((scalar_t *)vec.data(), {size}, c10::CppTypeToScalarType::value); return inplace ? out : out.clone(); } template inline c10::Dict from_vector(const phmap::flat_hash_map> &vec_dict, bool inplace = false) { c10::Dict out_dict; for (const auto &kv : vec_dict) out_dict.insert(kv.first, from_vector(kv.second, inplace)); return out_dict; } inline int64_t uniform_randint(int64_t low, int64_t high) { CHECK_LT(low, high); auto options = torch::TensorOptions().dtype(torch::kInt64); auto ret = torch::randint(low, high, {1}, options); auto ptr = ret.data_ptr(); return *ptr; } inline int64_t uniform_randint(int64_t high) { return uniform_randint(0, high); } inline torch::Tensor choice(int64_t population, int64_t num_samples, bool replace = false, torch::optional weight = torch::nullopt) { if (population == 0 || num_samples == 0) return torch::empty({0}, at::kLong); if (!replace && num_samples >= population) return torch::arange(population, at::kLong); if (weight.has_value()) return torch::multinomial(weight.value(), num_samples, replace); if (replace) { const auto out = torch::empty({num_samples}, at::kLong); auto *out_data = out.data_ptr(); for (int64_t i = 0; i < num_samples; i++) { out_data[i] = uniform_randint(population); } return out; } else { // Sample without replacement via Robert Floyd algorithm: // https://www.nowherenearithaca.com/2013/05/ // robert-floyds-tiny-and-beautiful.html const auto out = torch::empty({num_samples}, at::kLong); auto *out_data = out.data_ptr(); std::unordered_set samples; for (int64_t i = population - num_samples; i < population; i++) { int64_t sample = uniform_randint(i); if (!samples.insert(sample).second) { sample = i; samples.insert(sample); } out_data[i - population + num_samples] = sample; } return out; } } template inline void uniform_choice(const int64_t population, const int64_t num_samples, const int64_t *idx_data, std::vector *samples, phmap::flat_hash_map *to_local_node) { if (population == 0 || num_samples == 0) return; if (replace) { for (int64_t i = 0; i < num_samples; i++) { const int64_t &v = idx_data[uniform_randint(population)]; if (to_local_node->insert({v, samples->size()}).second) samples->push_back(v); } } else if (num_samples >= population) { for (int64_t i = 0; i < population; i++) { const int64_t &v = idx_data[i]; if (to_local_node->insert({v, samples->size()}).second) samples->push_back(v); } } else { std::unordered_set indices; for (int64_t i = population - num_samples; i < population; i++) { int64_t j = uniform_randint(i); if (!indices.insert(j).second) { j = i; indices.insert(j); } const int64_t &v = idx_data[j]; if (to_local_node->insert({v, samples->size()}).second) samples->push_back(v); } } } pytorch_sparse-0.6.18/csrc/cuda/000077500000000000000000000000001450761466100165335ustar00rootroot00000000000000pytorch_sparse-0.6.18/csrc/cuda/atomics.cuh000066400000000000000000000012241450761466100206720ustar00rootroot00000000000000#pragma once static inline __device__ void atomAdd(float *address, float val) { atomicAdd(address, val); } static inline __device__ void atomAdd(double *address, double val) { #if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)) unsigned long long int *address_as_ull = (unsigned long long int *)address; unsigned long long int old = *address_as_ull; unsigned long long int assumed; do { assumed = old; old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed))); } while (assumed != old); #else atomicAdd(address, val); #endif } pytorch_sparse-0.6.18/csrc/cuda/convert_cuda.cu000066400000000000000000000037701450761466100215470ustar00rootroot00000000000000#include "convert_cuda.h" #include #include "utils.cuh" #define THREADS 256 __global__ void ind2ptr_kernel(const int64_t *ind_data, int64_t *out_data, int64_t M, int64_t numel) { int64_t thread_idx = blockDim.x * blockIdx.x + threadIdx.x; if (thread_idx == 0) { for (int64_t i = 0; i <= ind_data[0]; i++) out_data[i] = 0; } else if (thread_idx < numel) { for (int64_t i = ind_data[thread_idx - 1]; i < ind_data[thread_idx]; i++) out_data[i + 1] = thread_idx; } else if (thread_idx == numel) { for (int64_t i = ind_data[numel - 1] + 1; i < M + 1; i++) out_data[i] = numel; } } torch::Tensor ind2ptr_cuda(torch::Tensor ind, int64_t M) { CHECK_CUDA(ind); cudaSetDevice(ind.get_device()); auto out = torch::empty({M + 1}, ind.options()); if (ind.numel() == 0) return out.zero_(); auto ind_data = ind.data_ptr(); auto out_data = out.data_ptr(); auto stream = at::cuda::getCurrentCUDAStream(); ind2ptr_kernel<<<(ind.numel() + 2 + THREADS - 1) / THREADS, THREADS, 0, stream>>>(ind_data, out_data, M, ind.numel()); return out; } __global__ void ptr2ind_kernel(const int64_t *ptr_data, int64_t *out_data, int64_t E, int64_t numel) { int64_t thread_idx = blockDim.x * blockIdx.x + threadIdx.x; if (thread_idx < numel) { int64_t idx = ptr_data[thread_idx], next_idx = ptr_data[thread_idx + 1]; for (int64_t i = idx; i < next_idx; i++) { out_data[i] = thread_idx; } } } torch::Tensor ptr2ind_cuda(torch::Tensor ptr, int64_t E) { CHECK_CUDA(ptr); cudaSetDevice(ptr.get_device()); auto out = torch::empty({E}, ptr.options()); auto ptr_data = ptr.data_ptr(); auto out_data = out.data_ptr(); auto stream = at::cuda::getCurrentCUDAStream(); ptr2ind_kernel<<<(ptr.numel() - 1 + THREADS - 1) / THREADS, THREADS, 0, stream>>>(ptr_data, out_data, E, ptr.numel() - 1); return out; } pytorch_sparse-0.6.18/csrc/cuda/convert_cuda.h000066400000000000000000000002361450761466100213610ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor ind2ptr_cuda(torch::Tensor ind, int64_t M); torch::Tensor ptr2ind_cuda(torch::Tensor ptr, int64_t E); pytorch_sparse-0.6.18/csrc/cuda/diag_cuda.cu000066400000000000000000000034731450761466100207730ustar00rootroot00000000000000#include "diag_cuda.h" #include #include "utils.cuh" #define THREADS 1024 __global__ void non_diag_mask_kernel(const int64_t *row_data, const int64_t *col_data, bool *out_data, int64_t N, int64_t k, int64_t num_diag, int64_t numel) { int64_t thread_idx = blockDim.x * blockIdx.x + threadIdx.x; if (thread_idx < numel) { int64_t r = row_data[thread_idx], c = col_data[thread_idx]; if (k < 0) { if (r + k < 0) { out_data[thread_idx] = true; } else if (r + k >= N) { out_data[thread_idx + num_diag] = true; } else if (r + k > c) { out_data[thread_idx + r + k] = true; } else if (r + k < c) { out_data[thread_idx + r + k + 1] = true; } } else { if (r + k >= N) { out_data[thread_idx + num_diag] = true; } else if (r + k > c) { out_data[thread_idx + r] = true; } else if (r + k < c) { out_data[thread_idx + r + 1] = true; } } } } torch::Tensor non_diag_mask_cuda(torch::Tensor row, torch::Tensor col, int64_t M, int64_t N, int64_t k) { CHECK_CUDA(row); CHECK_CUDA(col); cudaSetDevice(row.get_device()); auto E = row.size(0); auto num_diag = k < 0 ? std::min(M + k, N) : std::min(M, N - k); auto row_data = row.data_ptr(); auto col_data = col.data_ptr(); auto mask = torch::zeros({E + num_diag}, row.options().dtype(torch::kBool)); auto mask_data = mask.data_ptr(); if (E == 0) return mask; auto stream = at::cuda::getCurrentCUDAStream(); non_diag_mask_kernel<<<(E + THREADS - 1) / THREADS, THREADS, 0, stream>>>( row_data, col_data, mask_data, N, k, num_diag, E); return mask; } pytorch_sparse-0.6.18/csrc/cuda/diag_cuda.h000066400000000000000000000002641450761466100206060ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor non_diag_mask_cuda(torch::Tensor row, torch::Tensor col, int64_t M, int64_t N, int64_t k); pytorch_sparse-0.6.18/csrc/cuda/reducer.cuh000066400000000000000000000076351450761466100207000ustar00rootroot00000000000000#pragma once #include #include enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX }; const std::map reduce2REDUCE = { {"sum", SUM}, {"mean", MEAN}, {"mul", MUL}, {"div", DIV}, {"min", MIN}, {"max", MAX}, }; #define AT_DISPATCH_REDUCTION_TYPES(reduce, ...) \ [&] { \ switch (reduce2REDUCE.at(reduce)) { \ case SUM: { \ const ReductionType REDUCE = SUM; \ return __VA_ARGS__(); \ } \ case MEAN: { \ const ReductionType REDUCE = MEAN; \ return __VA_ARGS__(); \ } \ case MUL: { \ const ReductionType REDUCE = MUL; \ return __VA_ARGS__(); \ } \ case DIV: { \ const ReductionType REDUCE = DIV; \ return __VA_ARGS__(); \ } \ case MIN: { \ const ReductionType REDUCE = MIN; \ return __VA_ARGS__(); \ } \ case MAX: { \ const ReductionType REDUCE = MAX; \ return __VA_ARGS__(); \ } \ } \ }() template struct Reducer { static inline __host__ __device__ scalar_t init() { if (REDUCE == MUL || REDUCE == DIV) return (scalar_t)1; else if (REDUCE == MIN) return std::numeric_limits::max(); else if (REDUCE == MAX) return std::numeric_limits::lowest(); else return (scalar_t)0; } static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val, int64_t *arg, int64_t new_arg) { if (REDUCE == SUM || REDUCE == MEAN) *val = *val + new_val; else if (REDUCE == MUL) *val = *val * new_val; else if (REDUCE == DIV) *val = *val / new_val; else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) { *val = new_val; *arg = new_arg; } } static inline __host__ __device__ void write(scalar_t *address, scalar_t val, int64_t *arg_address, int64_t arg, int count) { if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV) *address = val; else if (REDUCE == MEAN) *address = val / (scalar_t)(count > 0 ? count : 1); else if (REDUCE == MIN || REDUCE == MAX) { if (count > 0) { *address = val; *arg_address = arg; } else *address = (scalar_t)0; } } }; pytorch_sparse-0.6.18/csrc/cuda/rw_cuda.cu000066400000000000000000000034761450761466100205220ustar00rootroot00000000000000#include "rw_cuda.h" #include #include "utils.cuh" #define THREADS 1024 #define BLOCKS(N) (N + THREADS - 1) / THREADS __global__ void uniform_random_walk_kernel(const int64_t *rowptr, const int64_t *col, const int64_t *start, const float *rand, int64_t *out, int64_t walk_length, int64_t numel) { const int64_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x; if (thread_idx < numel) { int64_t cur = start[thread_idx]; out[thread_idx] = cur; int64_t row_start, row_end; for (int64_t l = 0; l < walk_length; l++) { row_start = rowptr[cur], row_end = rowptr[cur + 1]; cur = col[row_start + int64_t(rand[l * numel + thread_idx] * (row_end - row_start))]; out[(l + 1) * numel + thread_idx] = cur; } } } torch::Tensor random_walk_cuda(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length) { CHECK_CUDA(rowptr); CHECK_CUDA(col); CHECK_CUDA(start); cudaSetDevice(rowptr.get_device()); CHECK_INPUT(rowptr.dim() == 1); CHECK_INPUT(col.dim() == 1); CHECK_INPUT(start.dim() == 1); auto rand = torch::rand({walk_length, start.size(0)}, start.options().dtype(torch::kFloat)); auto out = torch::full({walk_length + 1, start.size(0)}, -1, start.options()); auto stream = at::cuda::getCurrentCUDAStream(); uniform_random_walk_kernel<<>>( rowptr.data_ptr(), col.data_ptr(), start.data_ptr(), rand.data_ptr(), out.data_ptr(), walk_length, start.numel()); return out.t().contiguous(); } pytorch_sparse-0.6.18/csrc/cuda/rw_cuda.h000066400000000000000000000002741450761466100203330ustar00rootroot00000000000000#pragma once #include "../extensions.h" torch::Tensor random_walk_cuda(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length); pytorch_sparse-0.6.18/csrc/cuda/spmm_cuda.cu000066400000000000000000000171431450761466100210420ustar00rootroot00000000000000#include "spmm_cuda.h" #include #include "reducer.cuh" #include "utils.cuh" #define THREADS 256 #define FULL_MASK 0xffffffff // Paper: Design Principles for Sparse Matrix Multiplication on the GPU // Code: https://github.com/owensgroup/merge-spmm template __global__ void spmm_kernel(const int64_t *rowptr_data, const int64_t *col_data, const scalar_t *value_data, const scalar_t *mat_data, scalar_t *out_data, int64_t *arg_out_data, int B, int M, int N, int K) { // We ignore blockIdx.y here, because threads // across `blockIdx.y` are treated equally. int thread_idx = blockDim.x * blockIdx.x + threadIdx.x; int row = thread_idx >> 5; // thread_idx / 32 int lane_idx = thread_idx & (32 - 1); // thread_idx % 32 int batch_idx = row / M; // Compute the column index of `mat` in which the thread is operating. int mat_col_idx = lane_idx + (blockIdx.y << 5); // Compute the output index (row-major order). int out_idx = row * K + mat_col_idx; // Helper arrays for warp communication. int mat_row, mat_rows[32]; scalar_t val, vals[HAS_VALUE ? 32 : 1]; // Do not aggregate/write across the Y-axis (lane_idx < leftover). int leftover = K - (blockIdx.y << 5); if (batch_idx < B) { int row_start = __ldg(rowptr_data + (row % M)); int row_end = __ldg(rowptr_data + (row % M) + 1); int col_idx = row_start + lane_idx; scalar_t result = Reducer::init(); int64_t arg; // Iterate over all `col` indices in parallel within a warp. for (int c = row_start; c < row_end; c += 32) { if (col_idx < row_end) { // Coalesced memory access into `col` and `val`. mat_row = __ldg(col_data + col_idx) * K; if (HAS_VALUE) val = __ldg(value_data + col_idx); } else { mat_row = -1; if (HAS_VALUE) val = (scalar_t)0; } col_idx += 32; #pragma unroll for (int i = 0; i < 32; i++) { // Communication between all threads in a warp. mat_rows[i] = SHFL_SYNC(FULL_MASK, mat_row, i); if (HAS_VALUE) vals[i] = SHFL_SYNC(FULL_MASK, val, i); } #pragma unroll for (int i = 0; i < 32; i++) { if (lane_idx < leftover && mat_rows[i] != -1) { // Coalesced memory access into `mat`. val = __ldg(mat_data + batch_idx * N * K + mat_rows[i] + mat_col_idx); if (HAS_VALUE) val = vals[i] * val; Reducer::update(&result, val, &arg, c + i); } } } if (lane_idx < leftover) { // Coalesced write into `out`. Reducer::write(out_data + out_idx, result, arg_out_data + out_idx, arg, row_end - row_start); } } } std::tuple> spmm_cuda(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor mat, std::string reduce) { CHECK_CUDA(rowptr); CHECK_CUDA(col); if (optional_value.has_value()) CHECK_CUDA(optional_value.value()); CHECK_CUDA(mat); cudaSetDevice(rowptr.get_device()); CHECK_INPUT(rowptr.dim() == 1); CHECK_INPUT(col.dim() == 1); if (optional_value.has_value()) { CHECK_INPUT(optional_value.value().dim() == 1); CHECK_INPUT(optional_value.value().size(0) == col.size(0)); } CHECK_INPUT(mat.dim() >= 2); mat = mat.contiguous(); auto sizes = mat.sizes().vec(); sizes[mat.dim() - 2] = rowptr.numel() - 1; auto out = torch::empty(sizes, mat.options()); torch::optional arg_out = torch::nullopt; int64_t *arg_out_data = nullptr; if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) { arg_out = torch::full_like(out, col.numel(), rowptr.options()); arg_out_data = arg_out.value().data_ptr(); } auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto M = rowptr.numel() - 1; auto N = mat.size(-2); auto K = mat.size(-1); auto B = mat.numel() / (N * K); auto BLOCKS = dim3((32 * B * M + THREADS - 1) / THREADS, (K + 31) / 32); auto stream = at::cuda::getCurrentCUDAStream(); AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, mat.scalar_type(), "_", [&] { auto mat_data = mat.data_ptr(); auto out_data = out.data_ptr(); AT_DISPATCH_REDUCTION_TYPES(reduce, [&] { if (optional_value.has_value()) { auto value_data = optional_value.value().data_ptr(); spmm_kernel<<>>( rowptr_data, col_data, value_data, mat_data, out_data, arg_out_data, B, M, N, K); } else { spmm_kernel<<>>( rowptr_data, col_data, nullptr, mat_data, out_data, arg_out_data, B, M, N, K); } }); }); return std::make_tuple(out, arg_out); } template __global__ void spmm_value_bw_kernel(const int64_t *row_data, const int64_t *rowptr_data, const int64_t *col_data, const scalar_t *mat_data, const scalar_t *grad_data, scalar_t *out_data, int B, int M, int N, int E, int K) { int thread_idx = blockDim.x * blockIdx.x + threadIdx.x; int index_idx = (thread_idx >> 5); // thread_idx / 32 int lane_idx = thread_idx & (32 - 1); // thread_idx % 32 if (index_idx < E) { int row = __ldg(row_data + index_idx); int col = __ldg(col_data + index_idx); scalar_t val = (scalar_t)0; for (int b = 0; b < B; b++) { for (int k = lane_idx; k < K; k += 32) { val += mat_data[b * N * K + col * K + k] * grad_data[b * M * K + row * K + k]; } } #pragma unroll for (int i = 32 / 2; i > 0; i /= 2) { // Parallel reduction inside a warp. val += SHFL_DOWN_SYNC(FULL_MASK, val, i); } if (lane_idx == 0) { if (REDUCE == MEAN) { int row_start = __ldg(rowptr_data + row); int row_end = __ldg(rowptr_data + row + 1); val /= (scalar_t)max(row_end - row_start, 1); } out_data[index_idx] = val; } } } torch::Tensor spmm_value_bw_cuda(torch::Tensor row, torch::Tensor rowptr, torch::Tensor col, torch::Tensor mat, torch::Tensor grad, std::string reduce) { CHECK_CUDA(row); CHECK_CUDA(rowptr); CHECK_CUDA(col); CHECK_CUDA(mat); CHECK_CUDA(grad); cudaSetDevice(row.get_device()); mat = mat.contiguous(); grad = grad.contiguous(); auto M = grad.size(-2); auto N = mat.size(-2); auto E = row.numel(); auto K = mat.size(-1); auto B = mat.numel() / (N * K); auto BLOCKS = dim3((E * 32 + THREADS - 1) / THREADS); auto out = torch::zeros({row.numel()}, grad.options()); auto row_data = row.data_ptr(); auto rowptr_data = rowptr.data_ptr(); auto col_data = col.data_ptr(); auto stream = at::cuda::getCurrentCUDAStream(); AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, mat.scalar_type(), "_", [&] { auto mat_data = mat.data_ptr(); auto grad_data = grad.data_ptr(); auto out_data = out.data_ptr(); AT_DISPATCH_REDUCTION_TYPES(reduce, [&] { spmm_value_bw_kernel<<>>( row_data, rowptr_data, col_data, mat_data, grad_data, out_data, B, M, N, E, K); }); }); return out; } pytorch_sparse-0.6.18/csrc/cuda/spmm_cuda.h000066400000000000000000000007361450761466100206620ustar00rootroot00000000000000#pragma once #include "../extensions.h" std::tuple> spmm_cuda(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor mat, std::string reduce); torch::Tensor spmm_value_bw_cuda(torch::Tensor row, torch::Tensor rowptr, torch::Tensor col, torch::Tensor mat, torch::Tensor grad, std::string reduce); pytorch_sparse-0.6.18/csrc/cuda/utils.cuh000066400000000000000000000021401450761466100203710ustar00rootroot00000000000000#pragma once #include "../extensions.h" #define CHECK_CUDA(x) \ AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor") #define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch") __device__ __inline__ at::Half __shfl_sync(const unsigned mask, const at::Half var, const int srcLane) { return __shfl_sync(mask, var.operator __half(), srcLane); } __device__ __inline__ at::Half __shfl_down_sync(const unsigned mask, const at::Half var, const unsigned int delta) { return __shfl_down_sync(mask, var.operator __half(), delta); } #ifdef USE_ROCM __device__ __inline__ at::Half __ldg(const at::Half* ptr) { return __ldg(reinterpret_cast(ptr)); } #define SHFL_UP_SYNC(mask, var, delta) __shfl_up(var, delta) #define SHFL_DOWN_SYNC(mask, var, delta) __shfl_down(var, delta) #define SHFL_SYNC(mask, var, delta) __shfl(var, delta) #else #define SHFL_UP_SYNC __shfl_up_sync #define SHFL_DOWN_SYNC __shfl_down_sync #define SHFL_SYNC __shfl_sync #endif pytorch_sparse-0.6.18/csrc/diag.cpp000066400000000000000000000014571450761466100172360ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/diag_cpu.h" #ifdef WITH_CUDA #include "cuda/diag_cuda.h" #endif #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__diag_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__diag_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API torch::Tensor non_diag_mask(torch::Tensor row, torch::Tensor col, int64_t M, int64_t N, int64_t k) { if (row.device().is_cuda()) { #ifdef WITH_CUDA return non_diag_mask_cuda(row, col, M, N, k); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return non_diag_mask_cpu(row, col, M, N, k); } } static auto registry = torch::RegisterOperators().op( "torch_sparse::non_diag_mask", &non_diag_mask); pytorch_sparse-0.6.18/csrc/ego_sample.cpp000066400000000000000000000020511450761466100204340ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/ego_sample_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__ego_sample_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__ego_sample_cpu(void) { return NULL; } #endif #endif #endif // Returns `rowptr`, `col`, `n_id`, `e_id`, `ptr`, `root_n_id` SPARSE_API std::tuple ego_k_hop_sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t depth, int64_t num_neighbors, bool replace) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return ego_k_hop_sample_adj_cpu(rowptr, col, idx, depth, num_neighbors, replace); } } static auto registry = torch::RegisterOperators().op( "torch_sparse::ego_k_hop_sample_adj", &ego_k_hop_sample_adj); pytorch_sparse-0.6.18/csrc/extensions.h000066400000000000000000000000551450761466100201670ustar00rootroot00000000000000#include "macros.h" #include pytorch_sparse-0.6.18/csrc/hgt_sample.cpp000066400000000000000000000021051450761466100204440ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/hgt_sample_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__hgt_sample_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__hgt_sample_cpu(void) { return NULL; } #endif #endif #endif // Returns 'output_node_dict', 'row_dict', 'col_dict', 'output_edge_dict' SPARSE_API std::tuple, c10::Dict, c10::Dict, c10::Dict> hgt_sample(const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_samples_dict, const int64_t num_hops) { return hgt_sample_cpu(colptr_dict, row_dict, input_node_dict, num_samples_dict, num_hops); } static auto registry = torch::RegisterOperators().op("torch_sparse::hgt_sample", &hgt_sample); pytorch_sparse-0.6.18/csrc/macros.h000066400000000000000000000006661450761466100172640ustar00rootroot00000000000000#pragma once #ifdef _WIN32 #if defined(torchsparse_EXPORTS) #define SPARSE_API __declspec(dllexport) #else #define SPARSE_API __declspec(dllimport) #endif #else #define SPARSE_API #endif #if (defined __cpp_inline_variables) || __cplusplus >= 201703L #define SPARSE_INLINE_VARIABLE inline #else #ifdef _MSC_VER #define SPARSE_INLINE_VARIABLE __declspec(selectany) #else #define SPARSE_INLINE_VARIABLE __attribute__((weak)) #endif #endif pytorch_sparse-0.6.18/csrc/metis.cpp000066400000000000000000000044061450761466100174500ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/metis_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__metis_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, int64_t num_parts, bool recursive) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return partition_cpu(rowptr, col, optional_value, torch::nullopt, num_parts, recursive); } } SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return partition_cpu(rowptr, col, optional_value, optional_node_weight, num_parts, recursive); } } SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive, int64_t num_workers) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return mt_partition_cpu(rowptr, col, optional_value, optional_node_weight, num_parts, recursive, num_workers); } } static auto registry = torch::RegisterOperators() .op("torch_sparse::partition", &partition) .op("torch_sparse::partition2", &partition2) .op("torch_sparse::mt_partition", &mt_partition); pytorch_sparse-0.6.18/csrc/neighbor_sample.cpp000066400000000000000000000052231450761466100214630ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/neighbor_sample_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__neighbor_sample_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__neighbor_sample_cpu(void) { return NULL; } #endif #endif #endif // Returns 'output_node', 'row', 'col', 'output_edge' SPARSE_API std::tuple neighbor_sample(const torch::Tensor &colptr, const torch::Tensor &row, const torch::Tensor &input_node, const std::vector num_neighbors, const bool replace, const bool directed) { return neighbor_sample_cpu(colptr, row, input_node, num_neighbors, replace, directed); } SPARSE_API std::tuple, c10::Dict, c10::Dict, c10::Dict> hetero_neighbor_sample( const std::vector &node_types, const std::vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const int64_t num_hops, const bool replace, const bool directed) { return hetero_neighbor_sample_cpu( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, num_hops, replace, directed); } std::tuple, c10::Dict, c10::Dict, c10::Dict> hetero_temporal_neighbor_sample( const std::vector &node_types, const std::vector &edge_types, const c10::Dict &colptr_dict, const c10::Dict &row_dict, const c10::Dict &input_node_dict, const c10::Dict> &num_neighbors_dict, const c10::Dict &node_time_dict, const int64_t num_hops, const bool replace, const bool directed) { return hetero_temporal_neighbor_sample_cpu( node_types, edge_types, colptr_dict, row_dict, input_node_dict, num_neighbors_dict, node_time_dict, num_hops, replace, directed); } static auto registry = torch::RegisterOperators() .op("torch_sparse::neighbor_sample", &neighbor_sample) .op("torch_sparse::hetero_neighbor_sample", &hetero_neighbor_sample) .op("torch_sparse::hetero_temporal_neighbor_sample", &hetero_temporal_neighbor_sample); pytorch_sparse-0.6.18/csrc/relabel.cpp000066400000000000000000000025071450761466100177350ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/relabel_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__relabel_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__relabel_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API std::tuple relabel(torch::Tensor col, torch::Tensor idx) { if (col.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return relabel_cpu(col, idx); } } SPARSE_API std::tuple, torch::Tensor> relabel_one_hop(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor idx, bool bipartite) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return relabel_one_hop_cpu(rowptr, col, optional_value, idx, bipartite); } } static auto registry = torch::RegisterOperators() .op("torch_sparse::relabel", &relabel) .op("torch_sparse::relabel_one_hop", &relabel_one_hop); pytorch_sparse-0.6.18/csrc/rw.cpp000066400000000000000000000015051450761466100167540ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/rw_cpu.h" #ifdef WITH_CUDA #include "cuda/rw_cuda.h" #endif #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__rw_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA return random_walk_cuda(rowptr, col, start, walk_length); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return random_walk_cpu(rowptr, col, start, walk_length); } } static auto registry = torch::RegisterOperators().op("torch_sparse::random_walk", &random_walk); pytorch_sparse-0.6.18/csrc/saint.cpp000066400000000000000000000013771450761466100174510ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/saint_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__saint_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__saint_cpu(void) { return NULL; } #endif #endif #endif std::tuple subgraph(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row, torch::Tensor col) { if (idx.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return subgraph_cpu(idx, rowptr, row, col); } } static auto registry = torch::RegisterOperators().op("torch_sparse::saint_subgraph", &subgraph); pytorch_sparse-0.6.18/csrc/sample.cpp000066400000000000000000000015101450761466100176010ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/sample_cpu.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__sample_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__sample_cpu(void) { return NULL; } #endif #endif #endif SPARSE_API std::tuple sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t num_neighbors, bool replace) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA AT_ERROR("No CUDA version supported"); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return sample_adj_cpu(rowptr, col, idx, num_neighbors, replace); } } static auto registry = torch::RegisterOperators().op("torch_sparse::sample_adj", &sample_adj); pytorch_sparse-0.6.18/csrc/sparse.h000066400000000000000000000064471450761466100173000ustar00rootroot00000000000000#pragma once #include "extensions.h" #include "macros.h" namespace sparse { SPARSE_API int64_t cuda_version() noexcept; namespace detail { SPARSE_INLINE_VARIABLE int64_t _cuda_version = cuda_version(); } // namespace detail } // namespace sparse SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M); SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E); SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, int64_t num_parts, bool recursive); SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive); SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::optional optional_node_weight, int64_t num_parts, bool recursive, int64_t num_workers); SPARSE_API std::tuple relabel(torch::Tensor col, torch::Tensor idx); SPARSE_API std::tuple, torch::Tensor> relabel_one_hop(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor idx, bool bipartite); SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, int64_t walk_length); SPARSE_API std::tuple subgraph(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row, torch::Tensor col); SPARSE_API std::tuple sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, int64_t num_neighbors, bool replace); SPARSE_API torch::Tensor spmm_sum(torch::optional opt_row, torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::optional opt_colptr, torch::optional opt_csr2csc, torch::Tensor mat); SPARSE_API torch::Tensor spmm_mean(torch::optional opt_row, torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::optional opt_rowcount, torch::optional opt_colptr, torch::optional opt_csr2csc, torch::Tensor mat); SPARSE_API std::tuple spmm_min(torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::Tensor mat); SPARSE_API std::tuple spmm_max(torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::Tensor mat); pytorch_sparse-0.6.18/csrc/spmm.cpp000066400000000000000000000322121450761466100172770ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #include "cpu/spmm_cpu.h" #ifdef WITH_CUDA #include "cuda/spmm_cuda.h" #endif #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__spmm_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__spmm_cpu(void) { return NULL; } #endif #endif #endif std::tuple> spmm_fw(torch::Tensor rowptr, torch::Tensor col, torch::optional optional_value, torch::Tensor mat, std::string reduce) { if (rowptr.device().is_cuda()) { #ifdef WITH_CUDA return spmm_cuda(rowptr, col, optional_value, mat, reduce); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return spmm_cpu(rowptr, col, optional_value, mat, reduce); } } torch::Tensor spmm_value_bw(torch::Tensor row, torch::Tensor rowptr, torch::Tensor col, torch::Tensor mat, torch::Tensor grad, std::string reduce) { if (row.device().is_cuda()) { #ifdef WITH_CUDA return spmm_value_bw_cuda(row, rowptr, col, mat, grad, reduce); #else AT_ERROR("Not compiled with CUDA support"); #endif } else { return spmm_value_bw_cpu(row, rowptr, col, mat, grad, reduce); } } using torch::autograd::AutogradContext; using torch::autograd::Variable; using torch::autograd::variable_list; class SPMMSum : public torch::autograd::Function { public: static variable_list forward(AutogradContext *ctx, torch::optional opt_row, Variable rowptr, Variable col, Variable value, torch::optional opt_colptr, torch::optional opt_csr2csc, Variable mat, bool has_value) { if (has_value && torch::autograd::any_variable_requires_grad({value})) { AT_ASSERTM(opt_row.has_value(), "Argument `row` is missing"); } if (torch::autograd::any_variable_requires_grad({mat})) { AT_ASSERTM(opt_row.has_value(), "Argument `row` is missing"); AT_ASSERTM(opt_colptr.has_value(), "Argument `colptr` is missing"); AT_ASSERTM(opt_csr2csc.has_value(), "Argument `csr2csc` is missing"); } auto row = opt_row.has_value() ? opt_row.value() : col; auto colptr = opt_colptr.has_value() ? opt_colptr.value() : col; auto csr2csc = opt_csr2csc.has_value() ? opt_csr2csc.value() : col; torch::optional opt_value = torch::nullopt; if (has_value) opt_value = value; auto out = std::get<0>(spmm_fw(rowptr, col, opt_value, mat, "sum")); ctx->saved_data["has_value"] = has_value; ctx->save_for_backward({row, rowptr, col, value, colptr, csr2csc, mat}); return {out}; } static variable_list backward(AutogradContext *ctx, variable_list grad_outs) { auto has_value = ctx->saved_data["has_value"].toBool(); auto grad_out = grad_outs[0]; auto saved = ctx->get_saved_variables(); auto row = saved[0], rowptr = saved[1], col = saved[2], value = saved[3], colptr = saved[4], csr2csc = saved[5], mat = saved[6]; auto grad_value = Variable(); if (has_value > 0 && torch::autograd::any_variable_requires_grad({value})) { grad_value = spmm_value_bw(row, rowptr, col, mat, grad_out, "sum"); } auto grad_mat = Variable(); if (torch::autograd::any_variable_requires_grad({mat})) { torch::optional opt_value = torch::nullopt; if (has_value) opt_value = value.view({-1, 1}).index_select(0, csr2csc).view(-1); grad_mat = std::get<0>(spmm_fw(colptr, row.index_select(0, csr2csc), opt_value, grad_out, "sum")); } return {Variable(), Variable(), Variable(), grad_value, Variable(), Variable(), grad_mat, Variable()}; } }; class SPMMMean : public torch::autograd::Function { public: static variable_list forward(AutogradContext *ctx, torch::optional opt_row, Variable rowptr, Variable col, Variable value, torch::optional opt_rowcount, torch::optional opt_colptr, torch::optional opt_csr2csc, Variable mat, bool has_value) { if (has_value && torch::autograd::any_variable_requires_grad({value})) { AT_ASSERTM(opt_row.has_value(), "Argument `row` is missing"); } if (torch::autograd::any_variable_requires_grad({mat})) { AT_ASSERTM(opt_row.has_value(), "Argument `row` is missing"); AT_ASSERTM(opt_rowcount.has_value(), "Argument `rowcount` is missing"); AT_ASSERTM(opt_colptr.has_value(), "Argument `colptr` is missing"); AT_ASSERTM(opt_csr2csc.has_value(), "Argument `csr2csc` is missing"); } auto row = opt_row.has_value() ? opt_row.value() : col; auto rowcount = opt_rowcount.has_value() ? opt_rowcount.value() : col; auto colptr = opt_colptr.has_value() ? opt_colptr.value() : col; auto csr2csc = opt_csr2csc.has_value() ? opt_csr2csc.value() : col; torch::optional opt_value = torch::nullopt; if (has_value) opt_value = value; auto out = std::get<0>(spmm_fw(rowptr, col, opt_value, mat, "mean")); ctx->saved_data["has_value"] = has_value; ctx->save_for_backward( {row, rowptr, col, value, rowcount, colptr, csr2csc, mat}); return {out}; } static variable_list backward(AutogradContext *ctx, variable_list grad_outs) { auto has_value = ctx->saved_data["has_value"].toBool(); auto grad_out = grad_outs[0]; auto saved = ctx->get_saved_variables(); auto row = saved[0], rowptr = saved[1], col = saved[2], value = saved[3], rowcount = saved[4], colptr = saved[5], csr2csc = saved[6], mat = saved[7]; auto grad_value = Variable(); if (has_value > 0 && torch::autograd::any_variable_requires_grad({value})) { grad_value = spmm_value_bw(row, rowptr, col, mat, grad_out, "mean"); } auto grad_mat = Variable(); if (torch::autograd::any_variable_requires_grad({mat})) { row = row.index_select(0, csr2csc); rowcount = rowcount.index_select(0, row).toType(mat.scalar_type()); rowcount.masked_fill_(rowcount < 1, 1); if (has_value > 0) rowcount = value.view({-1, 1}).index_select(0, csr2csc).view(-1).div(rowcount); else rowcount.pow_(-1); grad_mat = std::get<0>(spmm_fw(colptr, row, rowcount, grad_out, "sum")); } return {Variable(), Variable(), Variable(), grad_value, Variable(), Variable(), Variable(), grad_mat, Variable()}; } }; class SPMMMin : public torch::autograd::Function { public: static variable_list forward(AutogradContext *ctx, Variable rowptr, Variable col, Variable value, Variable mat, bool has_value) { torch::optional opt_value = torch::nullopt; if (has_value) opt_value = value; auto result = spmm_fw(rowptr, col, opt_value, mat, "min"); auto out = std::get<0>(result); auto arg_out = std::get<1>(result).value(); ctx->saved_data["has_value"] = has_value; ctx->save_for_backward({col, value, mat, arg_out}); ctx->mark_non_differentiable({arg_out}); return {out, arg_out}; } static variable_list backward(AutogradContext *ctx, variable_list grad_outs) { auto has_value = ctx->saved_data["has_value"].toBool(); auto grad_out = grad_outs[0]; auto saved = ctx->get_saved_variables(); auto col = saved[0], value = saved[1], mat = saved[2], arg_out = saved[3]; auto invalid_arg_mask = arg_out == col.size(0); arg_out = arg_out.masked_fill(invalid_arg_mask, 0); auto grad_value = Variable(); if (has_value > 0 && torch::autograd::any_variable_requires_grad({value})) { auto ind = col.index_select(0, arg_out.flatten()).view_as(arg_out); auto out = mat.gather(-2, ind); out.mul_(grad_out); out.masked_fill_(invalid_arg_mask, 0); grad_value = torch::zeros_like(value); grad_value.scatter_add_(0, arg_out.flatten(), out.flatten()); } auto grad_mat = Variable(); if (torch::autograd::any_variable_requires_grad({mat})) { if (has_value > 0) { value = value.view({-1, 1}) .index_select(0, arg_out.flatten()) .view_as(arg_out) .mul_(grad_out); } else value = grad_out; value.masked_fill_(invalid_arg_mask, 0); auto ind = col.index_select(0, arg_out.flatten()).view_as(arg_out); grad_mat = torch::zeros_like(mat); grad_mat.scatter_add_(-2, ind, value); } return {Variable(), Variable(), grad_value, grad_mat, Variable()}; } }; class SPMMMax : public torch::autograd::Function { public: static variable_list forward(AutogradContext *ctx, Variable rowptr, Variable col, Variable value, Variable mat, bool has_value) { torch::optional opt_value = torch::nullopt; if (has_value) opt_value = value; auto result = spmm_fw(rowptr, col, opt_value, mat, "max"); auto out = std::get<0>(result); auto arg_out = std::get<1>(result).value(); ctx->saved_data["has_value"] = has_value; ctx->save_for_backward({col, value, mat, arg_out}); ctx->mark_non_differentiable({arg_out}); return {out, arg_out}; } static variable_list backward(AutogradContext *ctx, variable_list grad_outs) { auto has_value = ctx->saved_data["has_value"].toBool(); auto grad_out = grad_outs[0]; auto saved = ctx->get_saved_variables(); auto col = saved[0], value = saved[1], mat = saved[2], arg_out = saved[3]; auto invalid_arg_mask = arg_out == col.size(0); arg_out = arg_out.masked_fill(invalid_arg_mask, 0); auto grad_value = Variable(); if (has_value > 0 && torch::autograd::any_variable_requires_grad({value})) { auto ind = col.index_select(0, arg_out.flatten()).view_as(arg_out); auto out = mat.gather(-2, ind); out.mul_(grad_out); out.masked_fill_(invalid_arg_mask, 0); grad_value = torch::zeros_like(value); grad_value.scatter_add_(0, arg_out.flatten(), out.flatten()); } auto grad_mat = Variable(); if (torch::autograd::any_variable_requires_grad({mat})) { if (has_value > 0) { value = value.view({-1, 1}) .index_select(0, arg_out.flatten()) .view_as(arg_out) .mul_(grad_out); } else value = grad_out; value.masked_fill_(invalid_arg_mask, 0); auto ind = col.index_select(0, arg_out.flatten()).view_as(arg_out); grad_mat = torch::zeros_like(mat); grad_mat.scatter_add_(-2, ind, value); } return {Variable(), Variable(), grad_value, grad_mat, Variable()}; } }; SPARSE_API torch::Tensor spmm_sum(torch::optional opt_row, torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::optional opt_colptr, torch::optional opt_csr2csc, torch::Tensor mat) { auto value = opt_value.has_value() ? opt_value.value() : col; return SPMMSum::apply(opt_row, rowptr, col, value, opt_colptr, opt_csr2csc, mat, opt_value.has_value())[0]; } SPARSE_API torch::Tensor spmm_mean(torch::optional opt_row, torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::optional opt_rowcount, torch::optional opt_colptr, torch::optional opt_csr2csc, torch::Tensor mat) { auto value = opt_value.has_value() ? opt_value.value() : col; return SPMMMean::apply(opt_row, rowptr, col, value, opt_rowcount, opt_colptr, opt_csr2csc, mat, opt_value.has_value())[0]; } SPARSE_API std::tuple spmm_min(torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::Tensor mat) { auto value = opt_value.has_value() ? opt_value.value() : col; auto result = SPMMMin::apply(rowptr, col, value, mat, opt_value.has_value()); return std::make_tuple(result[0], result[1]); } SPARSE_API std::tuple spmm_max(torch::Tensor rowptr, torch::Tensor col, torch::optional opt_value, torch::Tensor mat) { auto value = opt_value.has_value() ? opt_value.value() : col; auto result = SPMMMax::apply(rowptr, col, value, mat, opt_value.has_value()); return std::make_tuple(result[0], result[1]); } static auto registry = torch::RegisterOperators() .op("torch_sparse::spmm_sum", &spmm_sum) .op("torch_sparse::spmm_mean", &spmm_mean) .op("torch_sparse::spmm_min", &spmm_min) .op("torch_sparse::spmm_max", &spmm_max); pytorch_sparse-0.6.18/csrc/version.cpp000066400000000000000000000013241450761466100200100ustar00rootroot00000000000000#ifdef WITH_PYTHON #include #endif #include #ifdef WITH_CUDA #ifdef USE_ROCM #include #else #include #endif #endif #include "macros.h" #ifdef _WIN32 #ifdef WITH_PYTHON #ifdef WITH_CUDA PyMODINIT_FUNC PyInit__version_cuda(void) { return NULL; } #else PyMODINIT_FUNC PyInit__version_cpu(void) { return NULL; } #endif #endif #endif namespace sparse { SPARSE_API int64_t cuda_version() noexcept { #ifdef WITH_CUDA #ifdef USE_ROCM return HIP_VERSION; #else return CUDA_VERSION; #endif #else return -1; #endif } } // namespace sparse static auto registry = torch::RegisterOperators().op( "torch_sparse::cuda_version", [] { return sparse::cuda_version(); }); pytorch_sparse-0.6.18/setup.cfg000066400000000000000000000011251450761466100165050ustar00rootroot00000000000000[metadata] long_description=file: README.md long_description_content_type=text/markdown classifiers = Development Status :: 5 - Production/Stable License :: OSI Approved :: MIT License Programming Language :: Python Programming Language :: Python :: 3.8 Programming Language :: Python :: 3.9 Programming Language :: Python :: 3.10 Programming Language :: Python :: 3.11 Programming Language :: Python :: 3 :: Only [aliases] test = pytest [tool:pytest] addopts = --capture=no [isort] multi_line_output=3 include_trailing_comma = True skip=.gitignore,__init__.py pytorch_sparse-0.6.18/setup.py000066400000000000000000000130131450761466100163750ustar00rootroot00000000000000import glob import os import os.path as osp import platform import sys from itertools import product import torch from setuptools import find_packages, setup from torch.__config__ import parallel_info from torch.utils.cpp_extension import ( CUDA_HOME, BuildExtension, CppExtension, CUDAExtension, ) __version__ = '0.6.18' URL = 'https://github.com/rusty1s/pytorch_sparse' WITH_CUDA = False if torch.cuda.is_available(): WITH_CUDA = CUDA_HOME is not None or torch.version.hip suffices = ['cpu', 'cuda'] if WITH_CUDA else ['cpu'] if os.getenv('FORCE_CUDA', '0') == '1': suffices = ['cuda', 'cpu'] if os.getenv('FORCE_ONLY_CUDA', '0') == '1': suffices = ['cuda'] if os.getenv('FORCE_ONLY_CPU', '0') == '1': suffices = ['cpu'] BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1' WITH_METIS = True if os.getenv('WITH_METIS', '0') == '1' else False WITH_MTMETIS = True if os.getenv('WITH_MTMETIS', '0') == '1' else False WITH_SYMBOLS = True if os.getenv('WITH_SYMBOLS', '0') == '1' else False def get_extensions(): extensions = [] extensions_dir = osp.join('csrc') main_files = glob.glob(osp.join(extensions_dir, '*.cpp')) # remove generated 'hip' files, in case of rebuilds main_files = [path for path in main_files if 'hip' not in path] for main, suffix in product(main_files, suffices): define_macros = [('WITH_PYTHON', None)] undef_macros = [] if sys.platform == 'win32': define_macros += [('torchsparse_EXPORTS', None)] libraries = [] if WITH_METIS: define_macros += [('WITH_METIS', None)] libraries += ['metis'] if WITH_MTMETIS: define_macros += [('WITH_MTMETIS', None)] define_macros += [('MTMETIS_64BIT_VERTICES', None)] define_macros += [('MTMETIS_64BIT_EDGES', None)] define_macros += [('MTMETIS_64BIT_WEIGHTS', None)] define_macros += [('MTMETIS_64BIT_PARTITIONS', None)] libraries += ['mtmetis', 'wildriver'] extra_compile_args = {'cxx': ['-O3']} if not os.name == 'nt': # Not on Windows: extra_compile_args['cxx'] += ['-Wno-sign-compare'] if sys.platform == 'darwin': # On macOS: extra_compile_args['cxx'] += ['-D_LIBCPP_DISABLE_AVAILABILITY'] extra_link_args = [] if WITH_SYMBOLS else ['-s'] info = parallel_info() if ('backend: OpenMP' in info and 'OpenMP not found' not in info and sys.platform != 'darwin'): extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP'] if sys.platform == 'win32': extra_compile_args['cxx'] += ['/openmp'] else: extra_compile_args['cxx'] += ['-fopenmp'] else: print('Compiling without OpenMP...') # Compile for mac arm64 if (sys.platform == 'darwin' and platform.machine() == 'arm64'): extra_compile_args['cxx'] += ['-arch', 'arm64'] extra_link_args += ['-arch', 'arm64'] if suffix == 'cuda': define_macros += [('WITH_CUDA', None)] nvcc_flags = os.getenv('NVCC_FLAGS', '') nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ') nvcc_flags += ['-O3'] if torch.version.hip: # USE_ROCM was added to later versions of PyTorch # Define here to support older PyTorch versions as well: define_macros += [('USE_ROCM', None)] undef_macros += ['__HIP_NO_HALF_CONVERSIONS__'] else: nvcc_flags += ['--expt-relaxed-constexpr'] extra_compile_args['nvcc'] = nvcc_flags name = main.split(os.sep)[-1][:-4] sources = [main] path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp') if osp.exists(path): sources += [path] path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu') if suffix == 'cuda' and osp.exists(path): sources += [path] phmap_dir = osp.abspath("third_party/parallel-hashmap") Extension = CppExtension if suffix == 'cpu' else CUDAExtension extension = Extension( f'torch_sparse._{name}_{suffix}', sources, include_dirs=[extensions_dir, phmap_dir], define_macros=define_macros, undef_macros=undef_macros, extra_compile_args=extra_compile_args, extra_link_args=extra_link_args, libraries=libraries, ) extensions += [extension] return extensions install_requires = [ 'scipy', ] test_requires = [ 'pytest', 'pytest-cov', ] # work-around hipify abs paths include_package_data = True if torch.cuda.is_available() and torch.version.hip: include_package_data = False setup( name='torch_sparse', version=__version__, description=('PyTorch Extension Library of Optimized Autograd Sparse ' 'Matrix Operations'), author='Matthias Fey', author_email='matthias.fey@tu-dortmund.de', url=URL, download_url=f'{URL}/archive/{__version__}.tar.gz', keywords=[ 'pytorch', 'sparse', 'sparse-matrices', 'autograd', ], python_requires='>=3.8', install_requires=install_requires, extras_require={ 'test': test_requires, }, ext_modules=get_extensions() if not BUILD_DOCS else [], cmdclass={ 'build_ext': BuildExtension.with_options(no_python_abi_suffix=True) }, packages=find_packages(), include_package_data=include_package_data, ) pytorch_sparse-0.6.18/test/000077500000000000000000000000001450761466100156445ustar00rootroot00000000000000pytorch_sparse-0.6.18/test/test_add.py000066400000000000000000000020201450761466100177770ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import SparseTensor, add from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_add(dtype, device): rowA = torch.tensor([0, 0, 1, 2, 2], device=device) colA = torch.tensor([0, 2, 1, 0, 1], device=device) valueA = tensor([1, 2, 4, 1, 3], dtype, device) A = SparseTensor(row=rowA, col=colA, value=valueA) rowB = torch.tensor([0, 0, 1, 2, 2], device=device) colB = torch.tensor([1, 2, 2, 1, 2], device=device) valueB = tensor([2, 3, 1, 2, 4], dtype, device) B = SparseTensor(row=rowB, col=colB, value=valueB) C = A + B rowC, colC, valueC = C.coo() assert rowC.tolist() == [0, 0, 0, 1, 1, 2, 2, 2] assert colC.tolist() == [0, 1, 2, 1, 2, 0, 1, 2] assert valueC.tolist() == [1, 2, 5, 4, 1, 1, 5, 4] @torch.jit.script def jit_add(A: SparseTensor, B: SparseTensor) -> SparseTensor: return add(A, B) jit_add(A, B) pytorch_sparse-0.6.18/test/test_cat.py000066400000000000000000000033511450761466100200260ustar00rootroot00000000000000import pytest import torch from torch_sparse.cat import cat from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices, tensor @pytest.mark.parametrize('device', devices) def test_cat(device): row, col = tensor([[0, 0, 1], [0, 1, 2]], torch.long, device) mat1 = SparseTensor(row=row, col=col) mat1.fill_cache_() row, col = tensor([[0, 0, 1, 2], [0, 1, 1, 0]], torch.long, device) mat2 = SparseTensor(row=row, col=col) mat2.fill_cache_() out = cat([mat1, mat2], dim=0) assert out.to_dense().tolist() == [[1, 1, 0], [0, 0, 1], [1, 1, 0], [0, 1, 0], [1, 0, 0]] assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.has_rowcount() assert out.storage.num_cached_keys() == 1 out = cat([mat1, mat2], dim=1) assert out.to_dense().tolist() == [[1, 1, 0, 1, 1], [0, 0, 1, 0, 1], [0, 0, 0, 1, 0]] assert out.storage.has_row() assert not out.storage.has_rowptr() assert out.storage.num_cached_keys() == 2 out = cat([mat1, mat2], dim=(0, 1)) assert out.to_dense().tolist() == [[1, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 1, 0]] assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.num_cached_keys() == 5 value = torch.randn((mat1.nnz(), 4), device=device) mat1 = mat1.set_value_(value, layout='coo') out = cat([mat1, mat1], dim=-1) assert out.storage.value().size() == (mat1.nnz(), 8) assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.num_cached_keys() == 5 pytorch_sparse-0.6.18/test/test_coalesce.py000066400000000000000000000021431450761466100210330ustar00rootroot00000000000000import torch from torch_sparse import coalesce def test_coalesce(): row = torch.tensor([1, 0, 1, 0, 2, 1]) col = torch.tensor([0, 1, 1, 1, 0, 0]) index = torch.stack([row, col], dim=0) index, _ = coalesce(index, None, m=3, n=2) assert index.tolist() == [[0, 1, 1, 2], [1, 0, 1, 0]] def test_coalesce_add(): row = torch.tensor([1, 0, 1, 0, 2, 1]) col = torch.tensor([0, 1, 1, 1, 0, 0]) index = torch.stack([row, col], dim=0) value = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]]) index, value = coalesce(index, value, m=3, n=2) assert index.tolist() == [[0, 1, 1, 2], [1, 0, 1, 0]] assert value.tolist() == [[6, 8], [7, 9], [3, 4], [5, 6]] def test_coalesce_max(): row = torch.tensor([1, 0, 1, 0, 2, 1]) col = torch.tensor([0, 1, 1, 1, 0, 0]) index = torch.stack([row, col], dim=0) value = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]]) index, value = coalesce(index, value, m=3, n=2, op='max') assert index.tolist() == [[0, 1, 1, 2], [1, 0, 1, 0]] assert value.tolist() == [[4, 5], [6, 7], [3, 4], [5, 6]] pytorch_sparse-0.6.18/test/test_convert.py000066400000000000000000000013151450761466100207350ustar00rootroot00000000000000import torch from torch_sparse import to_scipy, from_scipy from torch_sparse import to_torch_sparse, from_torch_sparse def test_convert_scipy(): index = torch.tensor([[0, 0, 1, 2, 2], [0, 2, 1, 0, 1]]) value = torch.Tensor([1, 2, 4, 1, 3]) N = 3 out = from_scipy(to_scipy(index, value, N, N)) assert out[0].tolist() == index.tolist() assert out[1].tolist() == value.tolist() def test_convert_torch_sparse(): index = torch.tensor([[0, 0, 1, 2, 2], [0, 2, 1, 0, 1]]) value = torch.Tensor([1, 2, 4, 1, 3]) N = 3 out = from_torch_sparse(to_torch_sparse(index, value, N, N).coalesce()) assert out[0].tolist() == index.tolist() assert out[1].tolist() == value.tolist() pytorch_sparse-0.6.18/test/test_diag.py000066400000000000000000000046521450761466100201700ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_remove_diag(dtype, device): row, col = tensor([[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device) value = tensor([1, 2, 3, 4], dtype, device) mat = SparseTensor(row=row, col=col, value=value) mat.fill_cache_() mat = mat.remove_diag() assert mat.storage.row().tolist() == [0, 1] assert mat.storage.col().tolist() == [1, 2] assert mat.storage.value().tolist() == [2, 3] assert mat.storage.num_cached_keys() == 2 assert mat.storage.rowcount().tolist() == [1, 1, 0] assert mat.storage.colcount().tolist() == [0, 1, 1] mat = SparseTensor(row=row, col=col, value=value) mat.fill_cache_() mat = mat.remove_diag(k=1) assert mat.storage.row().tolist() == [0, 2] assert mat.storage.col().tolist() == [0, 2] assert mat.storage.value().tolist() == [1, 4] assert mat.storage.num_cached_keys() == 2 assert mat.storage.rowcount().tolist() == [1, 0, 1] assert mat.storage.colcount().tolist() == [1, 0, 1] @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_set_diag(dtype, device): row, col = tensor([[0, 0, 9, 9], [0, 1, 0, 1]], torch.long, device) value = tensor([1, 2, 3, 4], dtype, device) mat = SparseTensor(row=row, col=col, value=value) mat = mat.set_diag(tensor([-8, -8], dtype, device), k=-1) mat = mat.set_diag(tensor([-8], dtype, device), k=1) @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_fill_diag(dtype, device): row, col = tensor([[0, 0, 9, 9], [0, 1, 0, 1]], torch.long, device) value = tensor([1, 2, 3, 4], dtype, device) mat = SparseTensor(row=row, col=col, value=value) mat = mat.fill_diag(-8, k=-1) mat = mat.fill_diag(-8, k=1) @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_get_diag(dtype, device): row, col = tensor([[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device) value = tensor([[1, 1], [2, 2], [3, 3], [4, 4]], dtype, device) mat = SparseTensor(row=row, col=col, value=value) assert mat.get_diag().tolist() == [[1, 1], [0, 0], [4, 4]] row, col = tensor([[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device) mat = SparseTensor(row=row, col=col) assert mat.get_diag().tolist() == [1, 0, 1] pytorch_sparse-0.6.18/test/test_ego_sample.py000066400000000000000000000016101450761466100213660ustar00rootroot00000000000000import torch from torch_sparse import SparseTensor def test_ego_k_hop_sample_adj(): rowptr = torch.tensor([0, 3, 5, 9, 10, 12, 14]) row = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 4, 4, 5, 5]) col = torch.tensor([1, 2, 3, 0, 2, 0, 1, 4, 5, 0, 2, 5, 2, 4]) _ = SparseTensor(row=row, col=col, sparse_sizes=(6, 6)) nid = torch.tensor([0, 1]) fn = torch.ops.torch_sparse.ego_k_hop_sample_adj out = fn(rowptr, col, nid, 1, 3, False) rowptr, col, nid, eid, ptr, root_n_id = out assert nid.tolist() == [0, 1, 2, 3, 0, 1, 2] assert rowptr.tolist() == [0, 3, 5, 7, 8, 10, 12, 14] # row [0, 0, 0, 1, 1, 2, 2, 3, 4, 4, 5, 5, 6, 6] assert col.tolist() == [1, 2, 3, 0, 2, 0, 1, 0, 5, 6, 4, 6, 4, 5] assert eid.tolist() == [0, 1, 2, 3, 4, 5, 6, 9, 0, 1, 3, 4, 5, 6] assert ptr.tolist() == [0, 4, 7] assert root_n_id.tolist() == [0, 5] pytorch_sparse-0.6.18/test/test_eye.py000066400000000000000000000044531450761466100200450ustar00rootroot00000000000000from itertools import product import pytest from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices, dtypes @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_eye(dtype, device): mat = SparseTensor.eye(3, dtype=dtype, device=device) assert mat.device() == device assert mat.storage.sparse_sizes() == (3, 3) assert mat.storage.row().tolist() == [0, 1, 2] assert mat.storage.rowptr().tolist() == [0, 1, 2, 3] assert mat.storage.col().tolist() == [0, 1, 2] assert mat.storage.value().tolist() == [1, 1, 1] assert mat.storage.value().dtype == dtype assert mat.storage.num_cached_keys() == 0 mat = SparseTensor.eye(3, has_value=False, device=device) assert mat.device() == device assert mat.storage.sparse_sizes() == (3, 3) assert mat.storage.row().tolist() == [0, 1, 2] assert mat.storage.rowptr().tolist() == [0, 1, 2, 3] assert mat.storage.col().tolist() == [0, 1, 2] assert mat.storage.value() is None assert mat.storage.num_cached_keys() == 0 mat = SparseTensor.eye(3, 4, fill_cache=True, device=device) assert mat.device() == device assert mat.storage.sparse_sizes() == (3, 4) assert mat.storage.row().tolist() == [0, 1, 2] assert mat.storage.rowptr().tolist() == [0, 1, 2, 3] assert mat.storage.col().tolist() == [0, 1, 2] assert mat.storage.num_cached_keys() == 5 assert mat.storage.rowcount().tolist() == [1, 1, 1] assert mat.storage.colptr().tolist() == [0, 1, 2, 3, 3] assert mat.storage.colcount().tolist() == [1, 1, 1, 0] assert mat.storage.csr2csc().tolist() == [0, 1, 2] assert mat.storage.csc2csr().tolist() == [0, 1, 2] mat = SparseTensor.eye(4, 3, fill_cache=True, device=device) assert mat.device() == device assert mat.storage.sparse_sizes() == (4, 3) assert mat.storage.row().tolist() == [0, 1, 2] assert mat.storage.rowptr().tolist() == [0, 1, 2, 3, 3] assert mat.storage.col().tolist() == [0, 1, 2] assert mat.storage.num_cached_keys() == 5 assert mat.storage.rowcount().tolist() == [1, 1, 1, 0] assert mat.storage.colptr().tolist() == [0, 1, 2, 3] assert mat.storage.colcount().tolist() == [1, 1, 1] assert mat.storage.csr2csc().tolist() == [0, 1, 2] assert mat.storage.csc2csr().tolist() == [0, 1, 2] pytorch_sparse-0.6.18/test/test_matmul.py000066400000000000000000000047121450761466100205600ustar00rootroot00000000000000from itertools import product import pytest import torch import torch_scatter from torch_sparse.matmul import matmul, spspmm from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices, grad_dtypes, reductions @pytest.mark.parametrize('dtype,device,reduce', product(grad_dtypes, devices, reductions)) def test_spmm(dtype, device, reduce): if device == torch.device('cuda:0') and dtype == torch.bfloat16: return # Not yet implemented. src = torch.randn((10, 8), dtype=dtype, device=device) src[2:4, :] = 0 # Remove multiple rows. src[:, 2:4] = 0 # Remove multiple columns. src = SparseTensor.from_dense(src).requires_grad_() row, col, value = src.coo() other = torch.randn((2, 8, 2), dtype=dtype, device=device, requires_grad=True) src_col = other.index_select(-2, col) * value.unsqueeze(-1) expected = torch_scatter.scatter(src_col, row, dim=-2, reduce=reduce) if reduce == 'min': expected[expected > 1000] = 0 if reduce == 'max': expected[expected < -1000] = 0 grad_out = torch.randn_like(expected) expected.backward(grad_out) expected_grad_value = value.grad value.grad = None expected_grad_other = other.grad other.grad = None out = matmul(src, other, reduce) out.backward(grad_out) atol = 1e-7 if dtype == torch.float16 or dtype == torch.bfloat16: atol = 1e-1 assert torch.allclose(expected, out, atol=atol) assert torch.allclose(expected_grad_value, value.grad, atol=atol) assert torch.allclose(expected_grad_other, other.grad, atol=atol) @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices)) def test_spspmm(dtype, device): if dtype in {torch.half, torch.bfloat16}: return # Not yet implemented. src = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=dtype, device=device) src = SparseTensor.from_dense(src) out = matmul(src, src) assert out.sizes() == [3, 3] assert out.has_value() rowptr, col, value = out.csr() assert rowptr.tolist() == [0, 1, 2, 3] assert col.tolist() == [0, 1, 2] assert value.tolist() == [1, 1, 1] src.set_value_(None) out = matmul(src, src) assert out.sizes() == [3, 3] assert not out.has_value() rowptr, col, value = out.csr() assert rowptr.tolist() == [0, 1, 2, 3] assert col.tolist() == [0, 1, 2] torch.jit.script(spspmm) pytorch_sparse-0.6.18/test/test_metis.py000066400000000000000000000024511450761466100204000ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices try: rowptr = torch.tensor([0, 1]) col = torch.tensor([0]) torch.ops.torch_sparse.partition(rowptr, col, None, 1, True) with_metis = True except RuntimeError: with_metis = False @pytest.mark.skipif(not with_metis, reason='Not compiled with METIS support') @pytest.mark.parametrize('device,weighted', product(devices, [False, True])) def test_metis(device, weighted): mat1 = torch.randn(6 * 6, device=device).view(6, 6) mat2 = torch.arange(6 * 6, dtype=torch.long, device=device).view(6, 6) mat3 = torch.ones(6 * 6, device=device).view(6, 6) vec1 = None vec2 = torch.rand(6, device=device) for mat, vec in product([mat1, mat2, mat3], [vec1, vec2]): mat = SparseTensor.from_dense(mat) _, partptr, perm = mat.partition(num_parts=1, recursive=False, weighted=weighted, node_weight=vec) assert partptr.numel() == 2 assert perm.numel() == 6 _, partptr, perm = mat.partition(num_parts=2, recursive=False, weighted=weighted, node_weight=vec) assert partptr.numel() == 3 assert perm.numel() == 6 pytorch_sparse-0.6.18/test/test_mul.py000066400000000000000000000031351450761466100200540ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import SparseTensor, mul from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_sparse_sparse_mul(dtype, device): rowA = torch.tensor([0, 0, 1, 2, 2], device=device) colA = torch.tensor([0, 2, 1, 0, 1], device=device) valueA = tensor([1, 2, 4, 1, 3], dtype, device) A = SparseTensor(row=rowA, col=colA, value=valueA) rowB = torch.tensor([0, 0, 1, 2, 2], device=device) colB = torch.tensor([1, 2, 2, 1, 2], device=device) valueB = tensor([2, 3, 1, 2, 4], dtype, device) B = SparseTensor(row=rowB, col=colB, value=valueB) C = A * B rowC, colC, valueC = C.coo() assert rowC.tolist() == [0, 2] assert colC.tolist() == [2, 1] assert valueC.tolist() == [6, 6] @torch.jit.script def jit_mul(A: SparseTensor, B: SparseTensor) -> SparseTensor: return mul(A, B) jit_mul(A, B) @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_sparse_sparse_mul_empty(dtype, device): rowA = torch.tensor([0], device=device) colA = torch.tensor([1], device=device) valueA = tensor([1], dtype, device) A = SparseTensor(row=rowA, col=colA, value=valueA) rowB = torch.tensor([1], device=device) colB = torch.tensor([0], device=device) valueB = tensor([2], dtype, device) B = SparseTensor(row=rowB, col=colB, value=valueB) C = A * B rowC, colC, valueC = C.coo() assert rowC.tolist() == [] assert colC.tolist() == [] assert valueC.tolist() == [] pytorch_sparse-0.6.18/test/test_neighbor_sample.py000066400000000000000000000025361450761466100224210ustar00rootroot00000000000000import torch from torch_sparse import SparseTensor neighbor_sample = torch.ops.torch_sparse.neighbor_sample def test_neighbor_sample(): adj = SparseTensor.from_edge_index(torch.tensor([[0], [1]])) colptr, row, _ = adj.csc() # Sampling in a non-directed way should not sample in wrong direction: out = neighbor_sample(colptr, row, torch.tensor([0]), [1], False, False) assert out[0].tolist() == [0] assert out[1].tolist() == [] assert out[2].tolist() == [] # Sampling should work: out = neighbor_sample(colptr, row, torch.tensor([1]), [1], False, False) assert out[0].tolist() == [1, 0] assert out[1].tolist() == [1] assert out[2].tolist() == [0] # Sampling with more hops: out = neighbor_sample(colptr, row, torch.tensor([1]), [1, 1], False, False) assert out[0].tolist() == [1, 0] assert out[1].tolist() == [1] assert out[2].tolist() == [0] def test_neighbor_sample_seed(): colptr = torch.tensor([0, 3, 6, 9]) row = torch.tensor([0, 1, 2, 0, 1, 2, 0, 1, 2]) input_nodes = torch.tensor([0, 1]) torch.manual_seed(42) out1 = neighbor_sample(colptr, row, input_nodes, [1, 1], True, False) torch.manual_seed(42) out2 = neighbor_sample(colptr, row, input_nodes, [1, 1], True, False) for data1, data2 in zip(out1, out2): assert data1.tolist() == data2.tolist() pytorch_sparse-0.6.18/test/test_overload.py000066400000000000000000000006521450761466100210730ustar00rootroot00000000000000import torch from torch_sparse.tensor import SparseTensor def test_overload(): row = torch.tensor([0, 1, 1, 2, 2]) col = torch.tensor([1, 0, 2, 1, 2]) mat = SparseTensor(row=row, col=col) other = torch.tensor([1, 2, 3]).view(3, 1) other + mat mat + other other * mat mat * other other = torch.tensor([1, 2, 3]).view(1, 3) other + mat mat + other other * mat mat * other pytorch_sparse-0.6.18/test/test_permute.py000066400000000000000000000011041450761466100207320ustar00rootroot00000000000000import pytest import torch from torch_sparse.tensor import SparseTensor from torch_sparse.testing import devices, tensor @pytest.mark.parametrize('device', devices) def test_permute(device): row, col = tensor([[0, 0, 1, 2, 2], [0, 1, 0, 1, 2]], torch.long, device) value = tensor([1, 2, 3, 4, 5], torch.float, device) adj = SparseTensor(row=row, col=col, value=value) row, col, value = adj.permute(torch.tensor([1, 0, 2])).coo() assert row.tolist() == [0, 1, 1, 2, 2] assert col.tolist() == [1, 0, 1, 0, 2] assert value.tolist() == [3, 2, 1, 4, 5] pytorch_sparse-0.6.18/test/test_saint.py000066400000000000000000000005111450761466100203700ustar00rootroot00000000000000import torch from torch_sparse.tensor import SparseTensor def test_saint_subgraph(): row = torch.tensor([0, 0, 1, 1, 2, 2, 2, 3, 3, 4]) col = torch.tensor([1, 2, 0, 2, 0, 1, 3, 2, 4, 3]) adj = SparseTensor(row=row, col=col) node_idx = torch.tensor([0, 1, 2]) adj, edge_index = adj.saint_subgraph(node_idx) pytorch_sparse-0.6.18/test/test_sample.py000066400000000000000000000023021450761466100205330ustar00rootroot00000000000000import torch from torch_sparse import SparseTensor, sample, sample_adj def test_sample(): row = torch.tensor([0, 0, 2, 2]) col = torch.tensor([1, 2, 0, 1]) adj = SparseTensor(row=row, col=col, sparse_sizes=(3, 3)) out = sample(adj, num_neighbors=1) assert out.min() >= 0 and out.max() <= 2 def test_sample_adj(): row = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 4, 4, 5, 5]) col = torch.tensor([1, 2, 3, 0, 2, 0, 1, 4, 5, 0, 2, 5, 2, 4]) value = torch.arange(row.size(0)) adj_t = SparseTensor(row=row, col=col, value=value, sparse_sizes=(6, 6)) out, n_id = sample_adj(adj_t, torch.arange(2, 6), num_neighbors=-1) assert n_id.tolist() == [2, 3, 4, 5, 0, 1] row, col, val = out.coo() assert row.tolist() == [0, 0, 0, 0, 1, 2, 2, 3, 3] assert col.tolist() == [2, 3, 4, 5, 4, 0, 3, 0, 2] assert val.tolist() == [7, 8, 5, 6, 9, 10, 11, 12, 13] out, n_id = sample_adj(adj_t, torch.arange(2, 6), num_neighbors=2, replace=True) assert out.nnz() == 8 out, n_id = sample_adj(adj_t, torch.arange(2, 6), num_neighbors=2, replace=False) assert out.nnz() == 7 # node 3 has only one edge... pytorch_sparse-0.6.18/test/test_spmm.py000066400000000000000000000011271450761466100202320ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import spmm from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_spmm(dtype, device): row = torch.tensor([0, 0, 1, 2, 2], device=device) col = torch.tensor([0, 2, 1, 0, 1], device=device) index = torch.stack([row, col], dim=0) value = tensor([1, 2, 4, 1, 3], dtype, device) x = tensor([[1, 4], [2, 5], [3, 6]], dtype, device) out = spmm(index, value, 3, 3, x) assert out.tolist() == [[7, 16], [8, 20], [7, 19]] pytorch_sparse-0.6.18/test/test_spspmm.py000066400000000000000000000033011450761466100205710ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import SparseTensor, spspmm from torch_sparse.testing import devices, grad_dtypes, tensor @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices)) def test_spspmm(dtype, device): if dtype in {torch.half, torch.bfloat16}: return # Not yet implemented. indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]], device=device) valueA = tensor([1, 2, 3, 4, 5], dtype, device) indexB = torch.tensor([[0, 2], [1, 0]], device=device) valueB = tensor([2, 4], dtype, device) indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2) assert indexC.tolist() == [[0, 1, 2], [0, 1, 1]] assert valueC.tolist() == [8, 6, 8] @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices)) def test_sparse_tensor_spspmm(dtype, device): if dtype in {torch.half, torch.bfloat16}: return # Not yet implemented. x = SparseTensor( row=torch.tensor( [0, 1, 1, 1, 2, 3, 4, 5, 5, 6, 6, 7, 7, 7, 8, 8, 9, 9], device=device), col=torch.tensor( [0, 5, 10, 15, 1, 2, 3, 7, 13, 6, 9, 5, 10, 15, 11, 14, 5, 15], device=device), value=torch.tensor([ 1, 3**-0.5, 3**-0.5, 3**-0.5, 1, 1, 1, -2**-0.5, -2**-0.5, -2**-0.5, -2**-0.5, 6**-0.5, -6**0.5 / 3, 6**-0.5, -2**-0.5, -2**-0.5, 2**-0.5, -2**-0.5 ], dtype=dtype, device=device), ) expected = torch.eye(10, device=device).to(dtype) out = x @ x.to_dense().t() assert torch.allclose(out, expected, atol=1e-2) out = x @ x.t() out = out.to_dense() assert torch.allclose(out, expected, atol=1e-2) pytorch_sparse-0.6.18/test/test_storage.py000066400000000000000000000136201450761466100207230ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse.storage import SparseStorage from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('device', devices) def test_ind2ptr(device): row = tensor([2, 2, 4, 5, 5, 6], torch.long, device) rowptr = torch.ops.torch_sparse.ind2ptr(row, 8) assert rowptr.tolist() == [0, 0, 0, 2, 2, 3, 5, 6, 6] row = torch.ops.torch_sparse.ptr2ind(rowptr, 6) assert row.tolist() == [2, 2, 4, 5, 5, 6] row = tensor([], torch.long, device) rowptr = torch.ops.torch_sparse.ind2ptr(row, 8) assert rowptr.tolist() == [0, 0, 0, 0, 0, 0, 0, 0, 0] row = torch.ops.torch_sparse.ptr2ind(rowptr, 0) assert row.tolist() == [] @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_storage(dtype, device): row, col = tensor([[0, 0, 1, 1], [0, 1, 0, 1]], torch.long, device) storage = SparseStorage(row=row, col=col) assert storage.row().tolist() == [0, 0, 1, 1] assert storage.col().tolist() == [0, 1, 0, 1] assert storage.value() is None assert storage.sparse_sizes() == (2, 2) row, col = tensor([[0, 0, 1, 1], [1, 0, 1, 0]], torch.long, device) value = tensor([2, 1, 4, 3], dtype, device) storage = SparseStorage(row=row, col=col, value=value) assert storage.row().tolist() == [0, 0, 1, 1] assert storage.col().tolist() == [0, 1, 0, 1] assert storage.value().tolist() == [1, 2, 3, 4] assert storage.sparse_sizes() == (2, 2) @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_caching(dtype, device): row, col = tensor([[0, 0, 1, 1], [0, 1, 0, 1]], torch.long, device) storage = SparseStorage(row=row, col=col) assert storage._row.tolist() == row.tolist() assert storage._col.tolist() == col.tolist() assert storage._value is None assert storage._rowcount is None assert storage._rowptr is None assert storage._colcount is None assert storage._colptr is None assert storage._csr2csc is None assert storage.num_cached_keys() == 0 storage.fill_cache_() assert storage._rowcount.tolist() == [2, 2] assert storage._rowptr.tolist() == [0, 2, 4] assert storage._colcount.tolist() == [2, 2] assert storage._colptr.tolist() == [0, 2, 4] assert storage._csr2csc.tolist() == [0, 2, 1, 3] assert storage._csc2csr.tolist() == [0, 2, 1, 3] assert storage.num_cached_keys() == 5 storage = SparseStorage(row=row, rowptr=storage._rowptr, col=col, value=storage._value, sparse_sizes=storage._sparse_sizes, rowcount=storage._rowcount, colptr=storage._colptr, colcount=storage._colcount, csr2csc=storage._csr2csc, csc2csr=storage._csc2csr) assert storage._rowcount.tolist() == [2, 2] assert storage._rowptr.tolist() == [0, 2, 4] assert storage._colcount.tolist() == [2, 2] assert storage._colptr.tolist() == [0, 2, 4] assert storage._csr2csc.tolist() == [0, 2, 1, 3] assert storage._csc2csr.tolist() == [0, 2, 1, 3] assert storage.num_cached_keys() == 5 storage.clear_cache_() assert storage._rowcount is None assert storage._rowptr is not None assert storage._colcount is None assert storage._colptr is None assert storage._csr2csc is None assert storage.num_cached_keys() == 0 @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_utility(dtype, device): row, col = tensor([[0, 0, 1, 1], [1, 0, 1, 0]], torch.long, device) value = tensor([1, 2, 3, 4], dtype, device) storage = SparseStorage(row=row, col=col, value=value) assert storage.has_value() storage.set_value_(value, layout='csc') assert storage.value().tolist() == [1, 3, 2, 4] storage.set_value_(value, layout='coo') assert storage.value().tolist() == [1, 2, 3, 4] storage = storage.set_value(value, layout='csc') assert storage.value().tolist() == [1, 3, 2, 4] storage = storage.set_value(value, layout='coo') assert storage.value().tolist() == [1, 2, 3, 4] storage = storage.sparse_resize((3, 3)) assert storage.sparse_sizes() == (3, 3) new_storage = storage.copy() assert new_storage != storage assert new_storage.col().data_ptr() == storage.col().data_ptr() new_storage = storage.clone() assert new_storage != storage assert new_storage.col().data_ptr() != storage.col().data_ptr() @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_coalesce(dtype, device): row, col = tensor([[0, 0, 0, 1, 1], [0, 1, 1, 0, 1]], torch.long, device) value = tensor([1, 1, 1, 3, 4], dtype, device) storage = SparseStorage(row=row, col=col, value=value) assert storage.row().tolist() == row.tolist() assert storage.col().tolist() == col.tolist() assert storage.value().tolist() == value.tolist() assert not storage.is_coalesced() storage = storage.coalesce() assert storage.is_coalesced() assert storage.row().tolist() == [0, 0, 1, 1] assert storage.col().tolist() == [0, 1, 0, 1] assert storage.value().tolist() == [1, 2, 3, 4] @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_sparse_reshape(dtype, device): row, col = tensor([[0, 1, 2, 3], [0, 1, 2, 3]], torch.long, device) storage = SparseStorage(row=row, col=col) storage = storage.sparse_reshape(2, 8) assert storage.sparse_sizes() == (2, 8) assert storage.row().tolist() == [0, 0, 1, 1] assert storage.col().tolist() == [0, 5, 2, 7] storage = storage.sparse_reshape(-1, 4) assert storage.sparse_sizes() == (4, 4) assert storage.row().tolist() == [0, 1, 2, 3] assert storage.col().tolist() == [0, 1, 2, 3] storage = storage.sparse_reshape(2, -1) assert storage.sparse_sizes() == (2, 8) assert storage.row().tolist() == [0, 0, 1, 1] assert storage.col().tolist() == [0, 5, 2, 7] pytorch_sparse-0.6.18/test/test_tensor.py000066400000000000000000000052331450761466100205720ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import SparseTensor from torch_sparse.testing import devices, grad_dtypes @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices)) def test_getitem(dtype, device): m = 50 n = 40 k = 10 mat = torch.randn(m, n, dtype=dtype, device=device) mat = SparseTensor.from_dense(mat) idx1 = torch.randint(0, m, (k, ), dtype=torch.long, device=device) idx2 = torch.randint(0, n, (k, ), dtype=torch.long, device=device) bool1 = torch.zeros(m, dtype=torch.bool, device=device) bool2 = torch.zeros(n, dtype=torch.bool, device=device) bool1.scatter_(0, idx1, 1) bool2.scatter_(0, idx2, 1) # idx1 and idx2 may have duplicates k1_bool = bool1.nonzero().size(0) k2_bool = bool2.nonzero().size(0) idx1np = idx1.cpu().numpy() idx2np = idx2.cpu().numpy() bool1np = bool1.cpu().numpy() bool2np = bool2.cpu().numpy() idx1list = idx1np.tolist() idx2list = idx2np.tolist() bool1list = bool1np.tolist() bool2list = bool2np.tolist() assert mat[:k, :k].sizes() == [k, k] assert mat[..., :k].sizes() == [m, k] assert mat[idx1, idx2].sizes() == [k, k] assert mat[idx1np, idx2np].sizes() == [k, k] assert mat[idx1list, idx2list].sizes() == [k, k] assert mat[bool1, bool2].sizes() == [k1_bool, k2_bool] assert mat[bool1np, bool2np].sizes() == [k1_bool, k2_bool] assert mat[bool1list, bool2list].sizes() == [k1_bool, k2_bool] assert mat[idx1].sizes() == [k, n] assert mat[idx1np].sizes() == [k, n] assert mat[idx1list].sizes() == [k, n] assert mat[bool1].sizes() == [k1_bool, n] assert mat[bool1np].sizes() == [k1_bool, n] assert mat[bool1list].sizes() == [k1_bool, n] @pytest.mark.parametrize('device', devices) def test_to_symmetric(device): row = torch.tensor([0, 0, 0, 1, 1], device=device) col = torch.tensor([0, 1, 2, 0, 2], device=device) value = torch.arange(1, 6, device=device) mat = SparseTensor(row=row, col=col, value=value) assert not mat.is_symmetric() mat = mat.to_symmetric() assert mat.is_symmetric() assert mat.to_dense().tolist() == [ [2, 6, 3], [6, 0, 5], [3, 5, 0], ] def test_equal(): row = torch.tensor([0, 0, 0, 1, 1]) col = torch.tensor([0, 1, 2, 0, 2]) value = torch.arange(1, 6) matA = SparseTensor(row=row, col=col, value=value) matB = SparseTensor(row=row, col=col, value=value) col = torch.tensor([0, 1, 2, 0, 1]) matC = SparseTensor(row=row, col=col, value=value) assert id(matA) != id(matB) assert matA == matB assert id(matA) != id(matC) assert matA != matC pytorch_sparse-0.6.18/test/test_transpose.py000066400000000000000000000021771450761466100213020ustar00rootroot00000000000000from itertools import product import pytest import torch from torch_sparse import transpose from torch_sparse.testing import devices, dtypes, tensor @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_transpose_matrix(dtype, device): row = torch.tensor([1, 0, 1, 2], device=device) col = torch.tensor([0, 1, 1, 0], device=device) index = torch.stack([row, col], dim=0) value = tensor([1, 2, 3, 4], dtype, device) index, value = transpose(index, value, m=3, n=2) assert index.tolist() == [[0, 0, 1, 1], [1, 2, 0, 1]] assert value.tolist() == [1, 4, 2, 3] @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) def test_transpose(dtype, device): row = torch.tensor([1, 0, 1, 0, 2, 1], device=device) col = torch.tensor([0, 1, 1, 1, 0, 0], device=device) index = torch.stack([row, col], dim=0) value = tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]], dtype, device) index, value = transpose(index, value, m=3, n=2) assert index.tolist() == [[0, 0, 1, 1], [1, 2, 0, 1]] assert value.tolist() == [[7, 9], [5, 6], [6, 8], [3, 4]] pytorch_sparse-0.6.18/third_party/000077500000000000000000000000001450761466100172165ustar00rootroot00000000000000pytorch_sparse-0.6.18/third_party/parallel-hashmap/000077500000000000000000000000001450761466100224315ustar00rootroot00000000000000pytorch_sparse-0.6.18/torch_sparse/000077500000000000000000000000001450761466100173615ustar00rootroot00000000000000pytorch_sparse-0.6.18/torch_sparse/__init__.py000066400000000000000000000065741450761466100215060ustar00rootroot00000000000000import importlib import os.path as osp import torch __version__ = '0.6.18' for library in [ '_version', '_convert', '_diag', '_spmm', '_metis', '_rw', '_saint', '_sample', '_ego_sample', '_hgt_sample', '_neighbor_sample', '_relabel' ]: cuda_spec = importlib.machinery.PathFinder().find_spec( f'{library}_cuda', [osp.dirname(__file__)]) cpu_spec = importlib.machinery.PathFinder().find_spec( f'{library}_cpu', [osp.dirname(__file__)]) spec = cuda_spec or cpu_spec if spec is not None: torch.ops.load_library(spec.origin) else: # pragma: no cover raise ImportError(f"Could not find module '{library}_cpu' in " f"{osp.dirname(__file__)}") cuda_version = torch.ops.torch_sparse.cuda_version() if torch.version.cuda is not None and cuda_version != -1: # pragma: no cover if cuda_version < 10000: major, minor = int(str(cuda_version)[0]), int(str(cuda_version)[2]) else: major, minor = int(str(cuda_version)[0:2]), int(str(cuda_version)[3]) t_major, t_minor = [int(x) for x in torch.version.cuda.split('.')] if t_major != major: raise RuntimeError( f'Detected that PyTorch and torch_sparse were compiled with ' f'different CUDA versions. PyTorch has CUDA version ' f'{t_major}.{t_minor} and torch_sparse has CUDA version ' f'{major}.{minor}. Please reinstall the torch_sparse that ' f'matches your PyTorch install.') from .storage import SparseStorage # noqa from .tensor import SparseTensor # noqa from .transpose import t # noqa from .narrow import narrow, __narrow_diag__ # noqa from .select import select # noqa from .index_select import index_select, index_select_nnz # noqa from .masked_select import masked_select, masked_select_nnz # noqa from .permute import permute # noqa from .diag import remove_diag, set_diag, fill_diag, get_diag # noqa from .add import add, add_, add_nnz, add_nnz_ # noqa from .mul import mul, mul_, mul_nnz, mul_nnz_ # noqa from .reduce import sum, mean, min, max # noqa from .matmul import matmul # noqa from .cat import cat # noqa from .rw import random_walk # noqa from .metis import partition # noqa from .bandwidth import reverse_cuthill_mckee # noqa from .saint import saint_subgraph # noqa from .sample import sample, sample_adj # noqa from .convert import to_torch_sparse, from_torch_sparse # noqa from .convert import to_scipy, from_scipy # noqa from .coalesce import coalesce # noqa from .transpose import transpose # noqa from .eye import eye # noqa from .spmm import spmm # noqa from .spspmm import spspmm # noqa from .spadd import spadd # noqa __all__ = [ 'SparseStorage', 'SparseTensor', 't', 'narrow', '__narrow_diag__', 'select', 'index_select', 'index_select_nnz', 'masked_select', 'masked_select_nnz', 'permute', 'remove_diag', 'set_diag', 'fill_diag', 'get_diag', 'add', 'add_', 'add_nnz', 'add_nnz_', 'mul', 'mul_', 'mul_nnz', 'mul_nnz_', 'sum', 'mean', 'min', 'max', 'matmul', 'cat', 'random_walk', 'partition', 'reverse_cuthill_mckee', 'saint_subgraph', 'to_torch_sparse', 'from_torch_sparse', 'to_scipy', 'from_scipy', 'coalesce', 'transpose', 'eye', 'spmm', 'spspmm', 'spadd', '__version__', ] pytorch_sparse-0.6.18/torch_sparse/add.py000066400000000000000000000067751450761466100205020ustar00rootroot00000000000000from typing import Optional import torch from torch import Tensor from torch_scatter import gather_csr from torch_sparse.tensor import SparseTensor @torch.jit._overload # noqa: F811 def add(src, other): # noqa: F811 # type: (SparseTensor, Tensor) -> SparseTensor pass @torch.jit._overload # noqa: F811 def add(src, other): # noqa: F811 # type: (SparseTensor, SparseTensor) -> SparseTensor pass def add(src, other): # noqa: F811 if isinstance(other, Tensor): rowptr, col, value = src.csr() if other.size(0) == src.size(0) and other.size(1) == 1: # Row-wise. other = gather_csr(other.squeeze(1), rowptr) elif other.size(0) == 1 and other.size(1) == src.size(1): # Col-wise. other = other.squeeze(0)[col] else: raise ValueError( f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or ' f'(1, {src.size(1)}, ...), but got size {other.size()}.') if value is not None: value = other.to(value.dtype).add_(value) else: value = other.add_(1) return src.set_value(value, layout='coo') elif isinstance(other, SparseTensor): rowA, colA, valueA = src.coo() rowB, colB, valueB = other.coo() row = torch.cat([rowA, rowB], dim=0) col = torch.cat([colA, colB], dim=0) value: Optional[Tensor] = None if valueA is not None and valueB is not None: value = torch.cat([valueA, valueB], dim=0) M = max(src.size(0), other.size(0)) N = max(src.size(1), other.size(1)) sparse_sizes = (M, N) out = SparseTensor(row=row, col=col, value=value, sparse_sizes=sparse_sizes) out = out.coalesce(reduce='sum') return out else: raise NotImplementedError def add_(src: SparseTensor, other: torch.Tensor) -> SparseTensor: rowptr, col, value = src.csr() if other.size(0) == src.size(0) and other.size(1) == 1: # Row-wise. other = gather_csr(other.squeeze(1), rowptr) elif other.size(0) == 1 and other.size(1) == src.size(1): # Col-wise. other = other.squeeze(0)[col] else: raise ValueError( f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or ' f'(1, {src.size(1)}, ...), but got size {other.size()}.') if value is not None: value = value.add_(other.to(value.dtype)) else: value = other.add_(1) return src.set_value_(value, layout='coo') def add_nnz(src: SparseTensor, other: torch.Tensor, layout: Optional[str] = None) -> SparseTensor: value = src.storage.value() if value is not None: value = value.add(other.to(value.dtype)) else: value = other.add(1) return src.set_value(value, layout=layout) def add_nnz_(src: SparseTensor, other: torch.Tensor, layout: Optional[str] = None) -> SparseTensor: value = src.storage.value() if value is not None: value = value.add_(other.to(value.dtype)) else: value = other.add(1) return src.set_value_(value, layout=layout) SparseTensor.add = lambda self, other: add(self, other) SparseTensor.add_ = lambda self, other: add_(self, other) SparseTensor.add_nnz = lambda self, other, layout=None: add_nnz( self, other, layout) SparseTensor.add_nnz_ = lambda self, other, layout=None: add_nnz_( self, other, layout) SparseTensor.__add__ = SparseTensor.add SparseTensor.__radd__ = SparseTensor.add SparseTensor.__iadd__ = SparseTensor.add_ pytorch_sparse-0.6.18/torch_sparse/bandwidth.py000066400000000000000000000013771450761466100217070ustar00rootroot00000000000000import scipy.sparse as sp from typing import Tuple, Optional import torch from torch_sparse.tensor import SparseTensor from torch_sparse.permute import permute def reverse_cuthill_mckee(src: SparseTensor, is_symmetric: Optional[bool] = None ) -> Tuple[SparseTensor, torch.Tensor]: if is_symmetric is None: is_symmetric = src.is_symmetric() if not is_symmetric: src = src.to_symmetric() sp_src = src.to_scipy(layout='csr') perm = sp.csgraph.reverse_cuthill_mckee(sp_src, symmetric_mode=True).copy() perm = torch.from_numpy(perm).to(torch.long).to(src.device()) out = permute(src, perm) return out, perm SparseTensor.reverse_cuthill_mckee = reverse_cuthill_mckee pytorch_sparse-0.6.18/torch_sparse/cat.py000066400000000000000000000200621450761466100205020ustar00rootroot00000000000000from typing import Optional, List, Tuple # noqa import torch from torch_sparse.storage import SparseStorage from torch_sparse.tensor import SparseTensor @torch.jit._overload # noqa: F811 def cat(tensors, dim): # noqa: F811 # type: (List[SparseTensor], int) -> SparseTensor pass @torch.jit._overload # noqa: F811 def cat(tensors, dim): # noqa: F811 # type: (List[SparseTensor], Tuple[int, int]) -> SparseTensor pass @torch.jit._overload # noqa: F811 def cat(tensors, dim): # noqa: F811 # type: (List[SparseTensor], List[int]) -> SparseTensor pass def cat(tensors, dim): # noqa: F811 assert len(tensors) > 0 if isinstance(dim, int): dim = tensors[0].dim() + dim if dim < 0 else dim if dim == 0: return cat_first(tensors) elif dim == 1: return cat_second(tensors) pass elif dim > 1 and dim < tensors[0].dim(): values = [] for tensor in tensors: value = tensor.storage.value() assert value is not None values.append(value) value = torch.cat(values, dim=dim - 1) return tensors[0].set_value(value, layout='coo') else: raise IndexError( (f'Dimension out of range: Expected to be in range of ' f'[{-tensors[0].dim()}, {tensors[0].dim() - 1}], but got ' f'{dim}.')) else: assert isinstance(dim, (tuple, list)) assert len(dim) == 2 assert sorted(dim) == [0, 1] return cat_diag(tensors) def cat_first(tensors: List[SparseTensor]) -> SparseTensor: rows: List[torch.Tensor] = [] rowptrs: List[torch.Tensor] = [] cols: List[torch.Tensor] = [] values: List[torch.Tensor] = [] sparse_sizes: List[int] = [0, 0] rowcounts: List[torch.Tensor] = [] nnz: int = 0 for tensor in tensors: row = tensor.storage._row if row is not None: rows.append(row + sparse_sizes[0]) rowptr = tensor.storage._rowptr if rowptr is not None: rowptrs.append(rowptr[1:] + nnz if len(rowptrs) > 0 else rowptr) cols.append(tensor.storage._col) value = tensor.storage._value if value is not None: values.append(value) rowcount = tensor.storage._rowcount if rowcount is not None: rowcounts.append(rowcount) sparse_sizes[0] += tensor.sparse_size(0) sparse_sizes[1] = max(sparse_sizes[1], tensor.sparse_size(1)) nnz += tensor.nnz() row: Optional[torch.Tensor] = None if len(rows) == len(tensors): row = torch.cat(rows, dim=0) rowptr: Optional[torch.Tensor] = None if len(rowptrs) == len(tensors): rowptr = torch.cat(rowptrs, dim=0) col = torch.cat(cols, dim=0) value: Optional[torch.Tensor] = None if len(values) == len(tensors): value = torch.cat(values, dim=0) rowcount: Optional[torch.Tensor] = None if len(rowcounts) == len(tensors): rowcount = torch.cat(rowcounts, dim=0) storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=(sparse_sizes[0], sparse_sizes[1]), rowcount=rowcount, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True) return tensors[0].from_storage(storage) def cat_second(tensors: List[SparseTensor]) -> SparseTensor: rows: List[torch.Tensor] = [] cols: List[torch.Tensor] = [] values: List[torch.Tensor] = [] sparse_sizes: List[int] = [0, 0] colptrs: List[torch.Tensor] = [] colcounts: List[torch.Tensor] = [] nnz: int = 0 for tensor in tensors: row, col, value = tensor.coo() rows.append(row) cols.append(tensor.storage._col + sparse_sizes[1]) if value is not None: values.append(value) colptr = tensor.storage._colptr if colptr is not None: colptrs.append(colptr[1:] + nnz if len(colptrs) > 0 else colptr) colcount = tensor.storage._colcount if colcount is not None: colcounts.append(colcount) sparse_sizes[0] = max(sparse_sizes[0], tensor.sparse_size(0)) sparse_sizes[1] += tensor.sparse_size(1) nnz += tensor.nnz() row = torch.cat(rows, dim=0) col = torch.cat(cols, dim=0) value: Optional[torch.Tensor] = None if len(values) == len(tensors): value = torch.cat(values, dim=0) colptr: Optional[torch.Tensor] = None if len(colptrs) == len(tensors): colptr = torch.cat(colptrs, dim=0) colcount: Optional[torch.Tensor] = None if len(colcounts) == len(tensors): colcount = torch.cat(colcounts, dim=0) storage = SparseStorage(row=row, rowptr=None, col=col, value=value, sparse_sizes=(sparse_sizes[0], sparse_sizes[1]), rowcount=None, colptr=colptr, colcount=colcount, csr2csc=None, csc2csr=None, is_sorted=False) return tensors[0].from_storage(storage) def cat_diag(tensors: List[SparseTensor]) -> SparseTensor: assert len(tensors) > 0 rows: List[torch.Tensor] = [] rowptrs: List[torch.Tensor] = [] cols: List[torch.Tensor] = [] values: List[torch.Tensor] = [] sparse_sizes: List[int] = [0, 0] rowcounts: List[torch.Tensor] = [] colptrs: List[torch.Tensor] = [] colcounts: List[torch.Tensor] = [] csr2cscs: List[torch.Tensor] = [] csc2csrs: List[torch.Tensor] = [] nnz: int = 0 for tensor in tensors: row = tensor.storage._row if row is not None: rows.append(row + sparse_sizes[0]) rowptr = tensor.storage._rowptr if rowptr is not None: rowptrs.append(rowptr[1:] + nnz if len(rowptrs) > 0 else rowptr) cols.append(tensor.storage._col + sparse_sizes[1]) value = tensor.storage._value if value is not None: values.append(value) rowcount = tensor.storage._rowcount if rowcount is not None: rowcounts.append(rowcount) colptr = tensor.storage._colptr if colptr is not None: colptrs.append(colptr[1:] + nnz if len(colptrs) > 0 else colptr) colcount = tensor.storage._colcount if colcount is not None: colcounts.append(colcount) csr2csc = tensor.storage._csr2csc if csr2csc is not None: csr2cscs.append(csr2csc + nnz) csc2csr = tensor.storage._csc2csr if csc2csr is not None: csc2csrs.append(csc2csr + nnz) sparse_sizes[0] += tensor.sparse_size(0) sparse_sizes[1] += tensor.sparse_size(1) nnz += tensor.nnz() row: Optional[torch.Tensor] = None if len(rows) == len(tensors): row = torch.cat(rows, dim=0) rowptr: Optional[torch.Tensor] = None if len(rowptrs) == len(tensors): rowptr = torch.cat(rowptrs, dim=0) col = torch.cat(cols, dim=0) value: Optional[torch.Tensor] = None if len(values) == len(tensors): value = torch.cat(values, dim=0) rowcount: Optional[torch.Tensor] = None if len(rowcounts) == len(tensors): rowcount = torch.cat(rowcounts, dim=0) colptr: Optional[torch.Tensor] = None if len(colptrs) == len(tensors): colptr = torch.cat(colptrs, dim=0) colcount: Optional[torch.Tensor] = None if len(colcounts) == len(tensors): colcount = torch.cat(colcounts, dim=0) csr2csc: Optional[torch.Tensor] = None if len(csr2cscs) == len(tensors): csr2csc = torch.cat(csr2cscs, dim=0) csc2csr: Optional[torch.Tensor] = None if len(csc2csrs) == len(tensors): csc2csr = torch.cat(csc2csrs, dim=0) storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=(sparse_sizes[0], sparse_sizes[1]), rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True) return tensors[0].from_storage(storage) pytorch_sparse-0.6.18/torch_sparse/coalesce.py000066400000000000000000000020041450761466100215050ustar00rootroot00000000000000import torch from torch_sparse.storage import SparseStorage def coalesce(index, value, m, n, op="add"): """Row-wise sorts :obj:`value` and removes duplicate entries. Duplicate entries are removed by scattering them together. For scattering, any operation of `"torch_scatter"`_ can be used. Args: index (:class:`LongTensor`): The index tensor of sparse matrix. value (:class:`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: :obj:`"add"`) :rtype: (:class:`LongTensor`, :class:`Tensor`) """ storage = SparseStorage(row=index[0], col=index[1], value=value, sparse_sizes=(m, n), is_sorted=False) storage = storage.coalesce(reduce=op) return torch.stack([storage.row(), storage.col()], dim=0), storage.value() pytorch_sparse-0.6.18/torch_sparse/convert.py000066400000000000000000000013231450761466100214120ustar00rootroot00000000000000import numpy as np import scipy.sparse import torch from torch import from_numpy def to_torch_sparse(index, value, m, n): return torch.sparse_coo_tensor(index.detach(), value, (m, n)) def from_torch_sparse(A): return A.indices().detach(), A.values() def to_scipy(index, value, m, n): assert not index.is_cuda and not value.is_cuda (row, col), data = index.detach(), value.detach() return scipy.sparse.coo_matrix((data, (row, col)), (m, n)) def from_scipy(A): A = A.tocoo() row, col, value = A.row.astype(np.int64), A.col.astype(np.int64), A.data row, col, value = from_numpy(row), from_numpy(col), from_numpy(value) index = torch.stack([row, col], dim=0) return index, value pytorch_sparse-0.6.18/torch_sparse/diag.py000066400000000000000000000075411450761466100206460ustar00rootroot00000000000000from typing import Optional import torch from torch import Tensor from torch_sparse.storage import SparseStorage from torch_sparse.tensor import SparseTensor def remove_diag(src: SparseTensor, k: int = 0) -> SparseTensor: row, col, value = src.coo() inv_mask = row != col if k == 0 else row != (col - k) new_row, new_col = row[inv_mask], col[inv_mask] if value is not None: value = value[inv_mask] rowcount = src.storage._rowcount colcount = src.storage._colcount if rowcount is not None or colcount is not None: mask = ~inv_mask if rowcount is not None: rowcount = rowcount.clone() rowcount[row[mask]] -= 1 if colcount is not None: colcount = colcount.clone() colcount[col[mask]] -= 1 storage = SparseStorage(row=new_row, rowptr=None, col=new_col, value=value, sparse_sizes=src.sparse_sizes(), rowcount=rowcount, colptr=None, colcount=colcount, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) def set_diag(src: SparseTensor, values: Optional[Tensor] = None, k: int = 0) -> SparseTensor: src = remove_diag(src, k=k) row, col, value = src.coo() mask = torch.ops.torch_sparse.non_diag_mask(row, col, src.size(0), src.size(1), k) inv_mask = ~mask start, num_diag = -k if k < 0 else 0, mask.numel() - row.numel() diag = torch.arange(start, start + num_diag, device=row.device) new_row = row.new_empty(mask.size(0)) new_row[mask] = row new_row[inv_mask] = diag new_col = col.new_empty(mask.size(0)) new_col[mask] = col new_col[inv_mask] = diag.add_(k) new_value: Optional[Tensor] = None if value is not None: new_value = value.new_empty((mask.size(0), ) + value.size()[1:]) new_value[mask] = value if values is not None: new_value[inv_mask] = values else: new_value[inv_mask] = torch.ones((num_diag, ), dtype=value.dtype, device=value.device) rowcount = src.storage._rowcount if rowcount is not None: rowcount = rowcount.clone() rowcount[start:start + num_diag] += 1 colcount = src.storage._colcount if colcount is not None: colcount = colcount.clone() colcount[start + k:start + num_diag + k] += 1 storage = SparseStorage(row=new_row, rowptr=None, col=new_col, value=new_value, sparse_sizes=src.sparse_sizes(), rowcount=rowcount, colptr=None, colcount=colcount, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) def fill_diag(src: SparseTensor, fill_value: float, k: int = 0) -> SparseTensor: num_diag = min(src.sparse_size(0), src.sparse_size(1) - k) if k < 0: num_diag = min(src.sparse_size(0) + k, src.sparse_size(1)) value = src.storage.value() if value is not None: sizes = [num_diag] + src.sizes()[2:] return set_diag(src, value.new_full(sizes, fill_value), k) else: return set_diag(src, None, k) def get_diag(src: SparseTensor) -> Tensor: row, col, value = src.coo() if value is None: value = torch.ones(row.size(0), device=row.device) sizes = list(value.size()) sizes[0] = min(src.size(0), src.size(1)) out = value.new_zeros(sizes) mask = row == col out[row[mask]] = value[mask] return out SparseTensor.remove_diag = lambda self, k=0: remove_diag(self, k) SparseTensor.set_diag = lambda self, values=None, k=0: set_diag( self, values, k) SparseTensor.fill_diag = lambda self, fill_value, k=0: fill_diag( self, fill_value, k) SparseTensor.get_diag = lambda self: get_diag(self) pytorch_sparse-0.6.18/torch_sparse/eye.py000066400000000000000000000013471450761466100205220ustar00rootroot00000000000000import torch def eye(m, dtype=None, device=None): """Returns a sparse matrix with ones on the diagonal and zeros elsewhere. Args: m (int): The first dimension of sparse matrix. dtype (`torch.dtype`, optional): The desired data type of returned value vector. (default is set by `torch.set_default_tensor_type()`) device (`torch.device`, optional): The desired device of returned tensors. (default is set by `torch.set_default_tensor_type()`) :rtype: (:class:`LongTensor`, :class:`Tensor`) """ row = torch.arange(m, dtype=torch.long, device=device) index = torch.stack([row, row], dim=0) value = torch.ones(m, dtype=dtype, device=device) return index, value pytorch_sparse-0.6.18/torch_sparse/index_select.py000066400000000000000000000064001450761466100224010ustar00rootroot00000000000000from typing import Optional import torch from torch_scatter import gather_csr from torch_sparse.storage import SparseStorage, get_layout from torch_sparse.tensor import SparseTensor def index_select(src: SparseTensor, dim: int, idx: torch.Tensor) -> SparseTensor: dim = src.dim() + dim if dim < 0 else dim assert idx.dim() == 1 if dim == 0: old_rowptr, col, value = src.csr() rowcount = src.storage.rowcount() rowcount = rowcount[idx] rowptr = col.new_zeros(idx.size(0) + 1) torch.cumsum(rowcount, dim=0, out=rowptr[1:]) row = torch.arange(idx.size(0), device=col.device).repeat_interleave(rowcount) perm = torch.arange(row.size(0), device=row.device) perm += gather_csr(old_rowptr[idx] - rowptr[:-1], rowptr) col = col[perm] if value is not None: value = value[perm] sparse_sizes = (idx.size(0), src.sparse_size(1)) storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=rowcount, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) elif dim == 1: old_colptr, row, value = src.csc() colcount = src.storage.colcount() colcount = colcount[idx] colptr = row.new_zeros(idx.size(0) + 1) torch.cumsum(colcount, dim=0, out=colptr[1:]) col = torch.arange(idx.size(0), device=row.device).repeat_interleave(colcount) perm = torch.arange(col.size(0), device=col.device) perm += gather_csr(old_colptr[idx] - colptr[:-1], colptr) row = row[perm] csc2csr = (idx.size(0) * row + col).argsort() row, col = row[csc2csr], col[csc2csr] if value is not None: value = value[perm][csc2csr] sparse_sizes = (src.sparse_size(0), idx.size(0)) storage = SparseStorage(row=row, rowptr=None, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=None, colptr=colptr, colcount=colcount, csr2csc=None, csc2csr=csc2csr, is_sorted=True) return src.from_storage(storage) else: value = src.storage.value() if value is not None: return src.set_value(value.index_select(dim - 1, idx), layout='coo') else: raise ValueError def index_select_nnz(src: SparseTensor, idx: torch.Tensor, layout: Optional[str] = None) -> SparseTensor: assert idx.dim() == 1 if get_layout(layout) == 'csc': idx = src.storage.csc2csr()[idx] row, col, value = src.coo() row, col = row[idx], col[idx] if value is not None: value = value[idx] return SparseTensor(row=row, rowptr=None, col=col, value=value, sparse_sizes=src.sparse_sizes(), is_sorted=True) SparseTensor.index_select = lambda self, dim, idx: index_select(self, dim, idx) tmp = lambda self, idx, layout=None: index_select_nnz( # noqa self, idx, layout) SparseTensor.index_select_nnz = tmp pytorch_sparse-0.6.18/torch_sparse/masked_select.py000066400000000000000000000060201450761466100225340ustar00rootroot00000000000000from typing import Optional import torch from torch_sparse.storage import SparseStorage, get_layout from torch_sparse.tensor import SparseTensor def masked_select(src: SparseTensor, dim: int, mask: torch.Tensor) -> SparseTensor: dim = src.dim() + dim if dim < 0 else dim assert mask.dim() == 1 storage = src.storage if dim == 0: row, col, value = src.coo() rowcount = src.storage.rowcount() rowcount = rowcount[mask] mask = mask[row] row = torch.arange(rowcount.size(0), device=row.device).repeat_interleave(rowcount) col = col[mask] if value is not None: value = value[mask] sparse_sizes = (rowcount.size(0), src.sparse_size(1)) storage = SparseStorage(row=row, rowptr=None, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=rowcount, colcount=None, colptr=None, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) elif dim == 1: row, col, value = src.coo() csr2csc = src.storage.csr2csc() row = row[csr2csc] col = col[csr2csc] colcount = src.storage.colcount() colcount = colcount[mask] mask = mask[col] col = torch.arange(colcount.size(0), device=col.device).repeat_interleave(colcount) row = row[mask] csc2csr = (colcount.size(0) * row + col).argsort() row, col = row[csc2csr], col[csc2csr] if value is not None: value = value[csr2csc][mask][csc2csr] sparse_sizes = (src.sparse_size(0), colcount.size(0)) storage = SparseStorage(row=row, rowptr=None, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=None, colcount=colcount, colptr=None, csr2csc=None, csc2csr=csc2csr, is_sorted=True) return src.from_storage(storage) else: value = src.storage.value() if value is not None: idx = mask.nonzero().flatten() return src.set_value(value.index_select(dim - 1, idx), layout='coo') else: raise ValueError def masked_select_nnz(src: SparseTensor, mask: torch.Tensor, layout: Optional[str] = None) -> SparseTensor: assert mask.dim() == 1 if get_layout(layout) == 'csc': mask = mask[src.storage.csc2csr()] row, col, value = src.coo() row, col = row[mask], col[mask] if value is not None: value = value[mask] return SparseTensor(row=row, rowptr=None, col=col, value=value, sparse_sizes=src.sparse_sizes(), is_sorted=True) SparseTensor.masked_select = lambda self, dim, mask: masked_select( self, dim, mask) tmp = lambda self, mask, layout=None: masked_select_nnz( # noqa self, mask, layout) SparseTensor.masked_select_nnz = tmp pytorch_sparse-0.6.18/torch_sparse/matmul.py000066400000000000000000000122421450761466100212330ustar00rootroot00000000000000from typing import Optional, Tuple import torch from torch import Tensor from torch_sparse.tensor import SparseTensor def spmm_sum(src: SparseTensor, other: torch.Tensor) -> torch.Tensor: rowptr, col, value = src.csr() row = src.storage._row csr2csc = src.storage._csr2csc colptr = src.storage._colptr if value is not None: value = value.to(other.dtype) if value is not None and value.requires_grad: row = src.storage.row() if other.requires_grad: row = src.storage.row() csr2csc = src.storage.csr2csc() colptr = src.storage.colptr() return torch.ops.torch_sparse.spmm_sum(row, rowptr, col, value, colptr, csr2csc, other) def spmm_add(src: SparseTensor, other: torch.Tensor) -> torch.Tensor: return spmm_sum(src, other) def spmm_mean(src: SparseTensor, other: torch.Tensor) -> torch.Tensor: rowptr, col, value = src.csr() row = src.storage._row rowcount = src.storage._rowcount csr2csc = src.storage._csr2csc colptr = src.storage._colptr if value is not None: value = value.to(other.dtype) if value is not None and value.requires_grad: row = src.storage.row() if other.requires_grad: row = src.storage.row() rowcount = src.storage.rowcount() csr2csc = src.storage.csr2csc() colptr = src.storage.colptr() return torch.ops.torch_sparse.spmm_mean(row, rowptr, col, value, rowcount, colptr, csr2csc, other) def spmm_min(src: SparseTensor, other: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: rowptr, col, value = src.csr() if value is not None: value = value.to(other.dtype) return torch.ops.torch_sparse.spmm_min(rowptr, col, value, other) def spmm_max(src: SparseTensor, other: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: rowptr, col, value = src.csr() if value is not None: value = value.to(other.dtype) return torch.ops.torch_sparse.spmm_max(rowptr, col, value, other) def spmm(src: SparseTensor, other: torch.Tensor, reduce: str = "sum") -> torch.Tensor: if reduce == 'sum' or reduce == 'add': return spmm_sum(src, other) elif reduce == 'mean': return spmm_mean(src, other) elif reduce == 'min': return spmm_min(src, other)[0] elif reduce == 'max': return spmm_max(src, other)[0] else: raise ValueError def spspmm_sum(src: SparseTensor, other: SparseTensor) -> SparseTensor: A = src.to_torch_sparse_coo_tensor() B = other.to_torch_sparse_coo_tensor() C = torch.sparse.mm(A, B) edge_index = C._indices() row, col = edge_index[0], edge_index[1] value: Optional[Tensor] = None if src.has_value() or other.has_value(): value = C._values() return SparseTensor( row=row, col=col, value=value, sparse_sizes=(C.size(0), C.size(1)), is_sorted=True, trust_data=True, ) def spspmm_add(src: SparseTensor, other: SparseTensor) -> SparseTensor: return spspmm_sum(src, other) def spspmm(src: SparseTensor, other: SparseTensor, reduce: str = "sum") -> SparseTensor: if reduce == 'sum' or reduce == 'add': return spspmm_sum(src, other) elif reduce == 'mean' or reduce == 'min' or reduce == 'max': raise NotImplementedError else: raise ValueError @torch.jit._overload # noqa: F811 def matmul(src, other, reduce): # noqa: F811 # type: (SparseTensor, torch.Tensor, str) -> torch.Tensor pass @torch.jit._overload # noqa: F811 def matmul(src, other, reduce): # noqa: F811 # type: (SparseTensor, SparseTensor, str) -> SparseTensor pass def matmul(src, other, reduce="sum"): # noqa: F811 """Matrix product of a sparse tensor with either another sparse tensor or a dense tensor. The sparse tensor represents an adjacency matrix and is stored as a list of edges. This method multiplies elements along the rows of the adjacency matrix with the column of the other matrix. In regular matrix multiplication, the products are then summed together, but this method allows us to use other aggregation functions as well. Args: src (:class:`SparseTensor`): The sparse tensor. other (:class:`Tensor` or :class:`SparseTensor`): The second matrix. reduce (string, optional): The function to reduce along the rows of :obj:`src` and columns of :obj:`other`. Can be :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`. (default: :obj:`"sum"`) :rtype: (:class:`Tensor`) """ if isinstance(other, torch.Tensor): return spmm(src, other, reduce) elif isinstance(other, SparseTensor): return spspmm(src, other, reduce) raise ValueError SparseTensor.spmm = lambda self, other, reduce="sum": spmm(self, other, reduce) SparseTensor.spspmm = lambda self, other, reduce="sum": spspmm( self, other, reduce) SparseTensor.matmul = lambda self, other, reduce="sum": matmul( self, other, reduce) SparseTensor.__matmul__ = lambda self, other: matmul(self, other, 'sum') pytorch_sparse-0.6.18/torch_sparse/metis.py000066400000000000000000000050641450761466100210610ustar00rootroot00000000000000from typing import Optional, Tuple import torch from torch import Tensor from torch_sparse.permute import permute from torch_sparse.tensor import SparseTensor def weight2metis(weight: Tensor) -> Optional[Tensor]: sorted_weight = weight.sort()[0] diff = sorted_weight[1:] - sorted_weight[:-1] if diff.sum() == 0: return None weight_min, weight_max = sorted_weight[0], sorted_weight[-1] srange = weight_max - weight_min min_diff = diff.min() scale = (min_diff / srange).item() tick, arange = scale.as_integer_ratio() weight_ratio = (weight - weight_min).div_(srange).mul_(arange).add_(tick) return weight_ratio.to(torch.long) def partition( src: SparseTensor, num_parts: int, recursive: bool = False, weighted: bool = False, node_weight: Optional[Tensor] = None, balance_edge: bool = False, ) -> Tuple[SparseTensor, Tensor, Tensor]: assert num_parts >= 1 if num_parts == 1: partptr = torch.tensor([0, src.size(0)], device=src.device()) perm = torch.arange(src.size(0), device=src.device()) return src, partptr, perm if balance_edge and node_weight: raise ValueError("Cannot set 'balance_edge' and 'node_weight' at the " "same time in 'torch_sparse.partition'") rowptr, col, value = src.csr() rowptr, col = rowptr.cpu(), col.cpu() if value is not None and weighted: assert value.numel() == col.numel() value = value.view(-1).detach().cpu() if value.is_floating_point(): value = weight2metis(value) else: value = None if balance_edge: node_weight = col.new_zeros(rowptr.numel() - 1) node_weight.scatter_add_(0, col, torch.ones_like(col)) if node_weight is not None: assert node_weight.numel() == rowptr.numel() - 1 node_weight = node_weight.view(-1).detach().cpu() if node_weight.is_floating_point(): node_weight = weight2metis(node_weight) cluster = torch.ops.torch_sparse.partition2(rowptr, col, value, node_weight, num_parts, recursive) else: cluster = torch.ops.torch_sparse.partition(rowptr, col, value, num_parts, recursive) cluster = cluster.to(src.device()) cluster, perm = cluster.sort() out = permute(src, perm) partptr = torch.ops.torch_sparse.ind2ptr(cluster, num_parts) return out, partptr, perm SparseTensor.partition = partition pytorch_sparse-0.6.18/torch_sparse/mul.py000066400000000000000000000077461450761466100205460ustar00rootroot00000000000000from typing import Optional import torch from torch import Tensor from torch_scatter import gather_csr from torch_sparse.tensor import SparseTensor @torch.jit._overload # noqa: F811 def mul(src, other): # noqa: F811 # type: (SparseTensor, Tensor) -> SparseTensor pass @torch.jit._overload # noqa: F811 def mul(src, other): # noqa: F811 # type: (SparseTensor, SparseTensor) -> SparseTensor pass def mul(src, other): # noqa: F811 if isinstance(other, Tensor): rowptr, col, value = src.csr() if other.size(0) == src.size(0) and other.size(1) == 1: # Row-wise... other = gather_csr(other.squeeze(1), rowptr) pass # Col-wise... elif other.size(0) == 1 and other.size(1) == src.size(1): other = other.squeeze(0)[col] else: raise ValueError( f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or ' f'(1, {src.size(1)}, ...), but got size {other.size()}.') if value is not None: value = other.to(value.dtype).mul_(value) else: value = other return src.set_value(value, layout='coo') assert isinstance(other, SparseTensor) if not src.is_coalesced(): raise ValueError("The `src` tensor is not coalesced") if not other.is_coalesced(): raise ValueError("The `other` tensor is not coalesced") rowA, colA, valueA = src.coo() rowB, colB, valueB = other.coo() row = torch.cat([rowA, rowB], dim=0) col = torch.cat([colA, colB], dim=0) if valueA is not None and valueB is not None: value = torch.cat([valueA, valueB], dim=0) else: raise ValueError('Both sparse tensors must contain values') M = max(src.size(0), other.size(0)) N = max(src.size(1), other.size(1)) sparse_sizes = (M, N) # Sort indices: idx = col.new_full((col.numel() + 1, ), -1) idx[1:] = row * sparse_sizes[1] + col perm = idx[1:].argsort() idx[1:] = idx[1:][perm] row, col, value = row[perm], col[perm], value[perm] valid_mask = idx[1:] == idx[:-1] valid_idx = valid_mask.nonzero().view(-1) return SparseTensor( row=row[valid_mask], col=col[valid_mask], value=value[valid_idx - 1] * value[valid_idx], sparse_sizes=sparse_sizes, ) def mul_(src: SparseTensor, other: torch.Tensor) -> SparseTensor: rowptr, col, value = src.csr() if other.size(0) == src.size(0) and other.size(1) == 1: # Row-wise... other = gather_csr(other.squeeze(1), rowptr) pass elif other.size(0) == 1 and other.size(1) == src.size(1): # Col-wise... other = other.squeeze(0)[col] else: raise ValueError( f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or ' f'(1, {src.size(1)}, ...), but got size {other.size()}.') if value is not None: value = value.mul_(other.to(value.dtype)) else: value = other return src.set_value_(value, layout='coo') def mul_nnz( src: SparseTensor, other: torch.Tensor, layout: Optional[str] = None, ) -> SparseTensor: value = src.storage.value() if value is not None: value = value.mul(other.to(value.dtype)) else: value = other return src.set_value(value, layout=layout) def mul_nnz_( src: SparseTensor, other: torch.Tensor, layout: Optional[str] = None, ) -> SparseTensor: value = src.storage.value() if value is not None: value = value.mul_(other.to(value.dtype)) else: value = other return src.set_value_(value, layout=layout) SparseTensor.mul = lambda self, other: mul(self, other) SparseTensor.mul_ = lambda self, other: mul_(self, other) SparseTensor.mul_nnz = lambda self, other, layout=None: mul_nnz( self, other, layout) SparseTensor.mul_nnz_ = lambda self, other, layout=None: mul_nnz_( self, other, layout) SparseTensor.__mul__ = SparseTensor.mul SparseTensor.__rmul__ = SparseTensor.mul SparseTensor.__imul__ = SparseTensor.mul_ pytorch_sparse-0.6.18/torch_sparse/narrow.py000066400000000000000000000110301450761466100212360ustar00rootroot00000000000000from typing import Tuple from torch_sparse.storage import SparseStorage from torch_sparse.tensor import SparseTensor def narrow(src: SparseTensor, dim: int, start: int, length: int) -> SparseTensor: if dim < 0: dim = src.dim() + dim if start < 0: start = src.size(dim) + start if dim == 0: rowptr, col, value = src.csr() rowptr = rowptr.narrow(0, start=start, length=length + 1) row_start = rowptr[0] rowptr = rowptr - row_start row_length = rowptr[-1] row = src.storage._row if row is not None: row = row.narrow(0, row_start, row_length) - start col = col.narrow(0, row_start, row_length) if value is not None: value = value.narrow(0, row_start, row_length) sparse_sizes = (length, src.sparse_size(1)) rowcount = src.storage._rowcount if rowcount is not None: rowcount = rowcount.narrow(0, start=start, length=length) storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=rowcount, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) elif dim == 1: # This is faster than accessing `csc()` contrary to the `dim=0` case. row, col, value = src.coo() mask = (col >= start) & (col < start + length) row = row[mask] col = col[mask] - start if value is not None: value = value[mask] sparse_sizes = (src.sparse_size(0), length) colptr = src.storage._colptr if colptr is not None: colptr = colptr.narrow(0, start=start, length=length + 1) colptr = colptr - colptr[0] colcount = src.storage._colcount if colcount is not None: colcount = colcount.narrow(0, start=start, length=length) storage = SparseStorage(row=row, rowptr=None, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=None, colptr=colptr, colcount=colcount, csr2csc=None, csc2csr=None, is_sorted=True) return src.from_storage(storage) else: value = src.storage.value() if value is not None: return src.set_value(value.narrow(dim - 1, start, length), layout='coo') else: raise ValueError def __narrow_diag__(src: SparseTensor, start: Tuple[int, int], length: Tuple[int, int]) -> SparseTensor: # This function builds the inverse operation of `cat_diag` and should hence # only be used on *diagonally stacked* sparse matrices. # That's the reason why this method is marked as *private*. rowptr, col, value = src.csr() rowptr = rowptr.narrow(0, start=start[0], length=length[0] + 1) row_start = int(rowptr[0]) rowptr = rowptr - row_start row_length = int(rowptr[-1]) row = src.storage._row if row is not None: row = row.narrow(0, row_start, row_length) - start[0] col = col.narrow(0, row_start, row_length) - start[1] if value is not None: value = value.narrow(0, row_start, row_length) sparse_sizes = length rowcount = src.storage._rowcount if rowcount is not None: rowcount = rowcount.narrow(0, start[0], length[0]) colptr = src.storage._colptr if colptr is not None: colptr = colptr.narrow(0, start[1], length[1] + 1) colptr = colptr - int(colptr[0]) # i.e. `row_start` colcount = src.storage._colcount if colcount is not None: colcount = colcount.narrow(0, start[1], length[1]) csr2csc = src.storage._csr2csc if csr2csc is not None: csr2csc = csr2csc.narrow(0, row_start, row_length) - row_start csc2csr = src.storage._csc2csr if csc2csr is not None: csc2csr = csc2csr.narrow(0, row_start, row_length) - row_start storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True) return src.from_storage(storage) SparseTensor.narrow = lambda self, dim, start, length: narrow( self, dim, start, length) SparseTensor.__narrow_diag__ = lambda self, start, length: __narrow_diag__( self, start, length) pytorch_sparse-0.6.18/torch_sparse/permute.py000066400000000000000000000004311450761466100214120ustar00rootroot00000000000000import torch from torch_sparse.tensor import SparseTensor def permute(src: SparseTensor, perm: torch.Tensor) -> SparseTensor: assert src.is_quadratic() return src.index_select(0, perm).index_select(1, perm) SparseTensor.permute = lambda self, perm: permute(self, perm) pytorch_sparse-0.6.18/torch_sparse/reduce.py000066400000000000000000000063151450761466100212070ustar00rootroot00000000000000from typing import Optional import torch from torch_scatter import scatter, segment_csr from torch_sparse.tensor import SparseTensor def reduction(src: SparseTensor, dim: Optional[int] = None, reduce: str = 'sum') -> torch.Tensor: value = src.storage.value() if dim is None: if value is not None: if reduce == 'sum' or reduce == 'add': return value.sum() elif reduce == 'mean': return value.mean() elif reduce == 'min': return value.min() elif reduce == 'max': return value.max() else: raise ValueError else: if reduce == 'sum' or reduce == 'add': return torch.tensor(src.nnz(), dtype=src.dtype(), device=src.device()) elif reduce == 'mean' or reduce == 'min' or reduce == 'max': return torch.tensor(1, dtype=src.dtype(), device=src.device()) else: raise ValueError else: if dim < 0: dim = src.dim() + dim if dim == 0 and value is not None: col = src.storage.col() return scatter(value, col, 0, None, src.size(1), reduce) elif dim == 0 and value is None: if reduce == 'sum' or reduce == 'add': return src.storage.colcount().to(src.dtype()) elif reduce == 'mean' or reduce == 'min' or reduce == 'max': return torch.ones(src.size(1), dtype=src.dtype()) else: raise ValueError elif dim == 1 and value is not None: return segment_csr(value, src.storage.rowptr(), None, reduce) elif dim == 1 and value is None: if reduce == 'sum' or reduce == 'add': return src.storage.rowcount().to(src.dtype()) elif reduce == 'mean' or reduce == 'min' or reduce == 'max': return torch.ones(src.size(0), dtype=src.dtype()) else: raise ValueError elif dim > 1 and value is not None: if reduce == 'sum' or reduce == 'add': return value.sum(dim=dim - 1) elif reduce == 'mean': return value.mean(dim=dim - 1) elif reduce == 'min': return value.min(dim=dim - 1)[0] elif reduce == 'max': return value.max(dim=dim - 1)[0] else: raise ValueError else: raise ValueError def sum(src: SparseTensor, dim: Optional[int] = None) -> torch.Tensor: return reduction(src, dim, reduce='sum') def mean(src: SparseTensor, dim: Optional[int] = None) -> torch.Tensor: return reduction(src, dim, reduce='mean') def min(src: SparseTensor, dim: Optional[int] = None) -> torch.Tensor: return reduction(src, dim, reduce='min') def max(src: SparseTensor, dim: Optional[int] = None) -> torch.Tensor: return reduction(src, dim, reduce='max') SparseTensor.sum = lambda self, dim=None: sum(self, dim) SparseTensor.mean = lambda self, dim=None: mean(self, dim) SparseTensor.min = lambda self, dim=None: min(self, dim) SparseTensor.max = lambda self, dim=None: max(self, dim) pytorch_sparse-0.6.18/torch_sparse/rw.py000066400000000000000000000004761450761466100203720ustar00rootroot00000000000000import torch from torch_sparse.tensor import SparseTensor def random_walk(src: SparseTensor, start: torch.Tensor, walk_length: int) -> torch.Tensor: rowptr, col, _ = src.csr() return torch.ops.torch_sparse.random_walk(rowptr, col, start, walk_length) SparseTensor.random_walk = random_walk pytorch_sparse-0.6.18/torch_sparse/saint.py000066400000000000000000000012731450761466100210540ustar00rootroot00000000000000from typing import Tuple import torch from torch_sparse.tensor import SparseTensor def saint_subgraph(src: SparseTensor, node_idx: torch.Tensor ) -> Tuple[SparseTensor, torch.Tensor]: row, col, value = src.coo() rowptr = src.storage.rowptr() data = torch.ops.torch_sparse.saint_subgraph(node_idx, rowptr, row, col) row, col, edge_index = data if value is not None: value = value[edge_index] out = SparseTensor(row=row, rowptr=None, col=col, value=value, sparse_sizes=(node_idx.size(0), node_idx.size(0)), is_sorted=True) return out, edge_index SparseTensor.saint_subgraph = saint_subgraph pytorch_sparse-0.6.18/torch_sparse/sample.py000066400000000000000000000023611450761466100212160ustar00rootroot00000000000000from typing import Optional, Tuple import torch from torch_sparse.tensor import SparseTensor def sample(src: SparseTensor, num_neighbors: int, subset: Optional[torch.Tensor] = None) -> torch.Tensor: rowptr, col, _ = src.csr() rowcount = src.storage.rowcount() if subset is not None: rowcount = rowcount[subset] rowptr = rowptr[subset] else: rowptr = rowptr[:-1] rand = torch.rand((rowcount.size(0), num_neighbors), device=col.device) rand.mul_(rowcount.to(rand.dtype).view(-1, 1)) rand = rand.to(torch.long) rand.add_(rowptr.view(-1, 1)) return col[rand] def sample_adj(src: SparseTensor, subset: torch.Tensor, num_neighbors: int, replace: bool = False) -> Tuple[SparseTensor, torch.Tensor]: rowptr, col, value = src.csr() rowptr, col, n_id, e_id = torch.ops.torch_sparse.sample_adj( rowptr, col, subset, num_neighbors, replace) if value is not None: value = value[e_id] out = SparseTensor(rowptr=rowptr, row=None, col=col, value=value, sparse_sizes=(subset.size(0), n_id.size(0)), is_sorted=True) return out, n_id SparseTensor.sample = sample SparseTensor.sample_adj = sample_adj pytorch_sparse-0.6.18/torch_sparse/select.py000066400000000000000000000004201450761466100212060ustar00rootroot00000000000000from torch_sparse.tensor import SparseTensor from torch_sparse.narrow import narrow def select(src: SparseTensor, dim: int, idx: int) -> SparseTensor: return narrow(src, dim, start=idx, length=1) SparseTensor.select = lambda self, dim, idx: select(self, dim, idx) pytorch_sparse-0.6.18/torch_sparse/spadd.py000066400000000000000000000013661450761466100210340ustar00rootroot00000000000000import torch from torch_sparse import coalesce def spadd(indexA, valueA, indexB, valueB, m, n): """Matrix addition of two sparse matrices. Args: indexA (:class:`LongTensor`): The index tensor of first sparse matrix. valueA (:class:`Tensor`): The value tensor of first sparse matrix. indexB (:class:`LongTensor`): The index tensor of second sparse matrix. valueB (:class:`Tensor`): The value tensor of second sparse matrix. m (int): The first dimension of the sparse matrices. n (int): The second dimension of the sparse matrices. """ index = torch.cat([indexA, indexB], dim=-1) value = torch.cat([valueA, valueB], dim=0) return coalesce(index=index, value=value, m=m, n=n, op='add') pytorch_sparse-0.6.18/torch_sparse/spmm.py000066400000000000000000000017571450761466100207210ustar00rootroot00000000000000from torch import Tensor from torch_scatter import scatter_add def spmm(index: Tensor, value: Tensor, m: int, n: int, matrix: Tensor) -> Tensor: """Matrix product of sparse matrix with dense matrix. Args: index (:class:`LongTensor`): The index tensor of sparse matrix. value (:class:`Tensor`): The value tensor of sparse matrix, either of floating-point or integer type. Does not work for boolean and complex number data types. m (int): The first dimension of sparse matrix. n (int): The second dimension of sparse matrix. matrix (:class:`Tensor`): The dense matrix of same type as :obj:`value`. :rtype: :class:`Tensor` """ assert n == matrix.size(-2) row, col = index[0], index[1] matrix = matrix if matrix.dim() > 1 else matrix.unsqueeze(-1) out = matrix.index_select(-2, col) out = out * value.unsqueeze(-1) out = scatter_add(out, row, dim=-2, dim_size=m) return out pytorch_sparse-0.6.18/torch_sparse/spspmm.py000066400000000000000000000026401450761466100212540ustar00rootroot00000000000000import torch from torch_sparse.tensor import SparseTensor from torch_sparse.matmul import matmul def spspmm(indexA, valueA, indexB, valueB, m, k, n, coalesced=False): """Matrix product of two sparse tensors. Both input sparse matrices need to be coalesced (use the :obj:`coalesced` attribute to force). Args: indexA (:class:`LongTensor`): The index tensor of first sparse matrix. valueA (:class:`Tensor`): The value tensor of first sparse matrix. indexB (:class:`LongTensor`): The index tensor of second sparse matrix. valueB (:class:`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 :obj:`True`, will coalesce both input sparse matrices. (default: :obj:`False`) :rtype: (:class:`LongTensor`, :class:`Tensor`) """ A = SparseTensor(row=indexA[0], col=indexA[1], value=valueA, sparse_sizes=(m, k), is_sorted=not coalesced) B = SparseTensor(row=indexB[0], col=indexB[1], value=valueB, sparse_sizes=(k, n), is_sorted=not coalesced) C = matmul(A, B) row, col, value = C.coo() return torch.stack([row, col], dim=0), value pytorch_sparse-0.6.18/torch_sparse/storage.py000066400000000000000000000610021450761466100213760ustar00rootroot00000000000000import warnings from typing import List, Optional, Tuple import torch from torch_scatter import scatter_add, segment_csr from torch_sparse.utils import Final, index_sort layouts: Final[List[str]] = ['coo', 'csr', 'csc'] def get_layout(layout: Optional[str] = None) -> str: if layout is None: layout = 'coo' warnings.warn('`layout` argument unset, using default layout ' '"coo". This may lead to unexpected behaviour.') assert layout == 'coo' or layout == 'csr' or layout == 'csc' return layout @torch.jit.script class SparseStorage(object): _row: Optional[torch.Tensor] _rowptr: Optional[torch.Tensor] _col: torch.Tensor _value: Optional[torch.Tensor] _sparse_sizes: Tuple[int, int] _rowcount: Optional[torch.Tensor] _colptr: Optional[torch.Tensor] _colcount: Optional[torch.Tensor] _csr2csc: Optional[torch.Tensor] _csc2csr: Optional[torch.Tensor] def __init__( self, row: Optional[torch.Tensor] = None, rowptr: Optional[torch.Tensor] = None, col: Optional[torch.Tensor] = None, value: Optional[torch.Tensor] = None, sparse_sizes: Optional[Tuple[Optional[int], Optional[int]]] = None, rowcount: Optional[torch.Tensor] = None, colptr: Optional[torch.Tensor] = None, colcount: Optional[torch.Tensor] = None, csr2csc: Optional[torch.Tensor] = None, csc2csr: Optional[torch.Tensor] = None, is_sorted: bool = False, trust_data: bool = False, ): assert row is not None or rowptr is not None assert col is not None assert col.dtype == torch.long assert col.dim() == 1 col = col.contiguous() M: int = 0 if sparse_sizes is None or sparse_sizes[0] is None: if rowptr is not None: M = rowptr.numel() - 1 elif row is not None and row.numel() > 0: M = int(row.max()) + 1 else: _M = sparse_sizes[0] assert _M is not None M = _M if rowptr is not None: assert rowptr.numel() - 1 == M elif row is not None and row.numel() > 0: assert trust_data or int(row.max()) < M N: int = 0 if sparse_sizes is None or sparse_sizes[1] is None: if col.numel() > 0: N = int(col.max()) + 1 else: _N = sparse_sizes[1] assert _N is not None N = _N if col.numel() > 0: assert trust_data or int(col.max()) < N sparse_sizes = (M, N) if row is not None: assert row.dtype == torch.long assert row.device == col.device assert row.dim() == 1 assert row.numel() == col.numel() row = row.contiguous() if rowptr is not None: assert rowptr.dtype == torch.long assert rowptr.device == col.device assert rowptr.dim() == 1 assert rowptr.numel() - 1 == sparse_sizes[0] rowptr = rowptr.contiguous() if value is not None: assert value.device == col.device assert value.size(0) == col.size(0) value = value.contiguous() if rowcount is not None: assert rowcount.dtype == torch.long assert rowcount.device == col.device assert rowcount.dim() == 1 assert rowcount.numel() == sparse_sizes[0] rowcount = rowcount.contiguous() if colptr is not None: assert colptr.dtype == torch.long assert colptr.device == col.device assert colptr.dim() == 1 assert colptr.numel() - 1 == sparse_sizes[1] colptr = colptr.contiguous() if colcount is not None: assert colcount.dtype == torch.long assert colcount.device == col.device assert colcount.dim() == 1 assert colcount.numel() == sparse_sizes[1] colcount = colcount.contiguous() if csr2csc is not None: assert csr2csc.dtype == torch.long assert csr2csc.device == col.device assert csr2csc.dim() == 1 assert csr2csc.numel() == col.size(0) csr2csc = csr2csc.contiguous() if csc2csr is not None: assert csc2csr.dtype == torch.long assert csc2csr.device == col.device assert csc2csr.dim() == 1 assert csc2csr.numel() == col.size(0) csc2csr = csc2csr.contiguous() self._row = row self._rowptr = rowptr self._col = col self._value = value self._sparse_sizes = tuple(sparse_sizes) self._rowcount = rowcount self._colptr = colptr self._colcount = colcount self._csr2csc = csr2csc self._csc2csr = csc2csr if not is_sorted and self._col.numel() > 0: idx = self._col.new_zeros(self._col.numel() + 1) idx[1:] = self.row() idx[1:] *= self._sparse_sizes[1] idx[1:] += self._col if (idx[1:] < idx[:-1]).any(): max_value = self._sparse_sizes[0] * self._sparse_sizes[1] _, perm = index_sort(idx[1:], max_value) self._row = self.row()[perm] self._col = self._col[perm] if value is not None: self._value = value[perm] self._csr2csc = None self._csc2csr = None @classmethod def empty(self): row = torch.tensor([], dtype=torch.long) col = torch.tensor([], dtype=torch.long) return SparseStorage( row=row, rowptr=None, col=col, value=None, sparse_sizes=(0, 0), rowcount=None, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True, trust_data=True, ) def has_row(self) -> bool: return self._row is not None def row(self): row = self._row if row is not None: return row rowptr = self._rowptr if rowptr is not None: row = torch.ops.torch_sparse.ptr2ind(rowptr, self._col.numel()) self._row = row return row raise ValueError def has_rowptr(self) -> bool: return self._rowptr is not None def rowptr(self) -> torch.Tensor: rowptr = self._rowptr if rowptr is not None: return rowptr row = self._row if row is not None: rowptr = torch.ops.torch_sparse.ind2ptr(row, self._sparse_sizes[0]) self._rowptr = rowptr return rowptr raise ValueError def col(self) -> torch.Tensor: return self._col def has_value(self) -> bool: return self._value is not None def value(self) -> Optional[torch.Tensor]: return self._value def set_value_( self, value: Optional[torch.Tensor], layout: Optional[str] = None, ): if value is not None: if get_layout(layout) == 'csc': value = value[self.csc2csr()] value = value.contiguous() assert value.device == self._col.device assert value.size(0) == self._col.numel() self._value = value return self def set_value( self, value: Optional[torch.Tensor], layout: Optional[str] = None, ): if value is not None: if get_layout(layout) == 'csc': value = value[self.csc2csr()] value = value.contiguous() assert value.device == self._col.device assert value.size(0) == self._col.numel() return SparseStorage( row=self._row, rowptr=self._rowptr, col=self._col, value=value, sparse_sizes=self._sparse_sizes, rowcount=self._rowcount, colptr=self._colptr, colcount=self._colcount, csr2csc=self._csr2csc, csc2csr=self._csc2csr, is_sorted=True, trust_data=True, ) def sparse_sizes(self) -> Tuple[int, int]: return self._sparse_sizes def sparse_size(self, dim: int) -> int: return self._sparse_sizes[dim] def sparse_resize(self, sparse_sizes: Tuple[int, int]): assert len(sparse_sizes) == 2 old_sparse_sizes, nnz = self._sparse_sizes, self._col.numel() diff_0 = sparse_sizes[0] - old_sparse_sizes[0] rowcount, rowptr = self._rowcount, self._rowptr if diff_0 > 0: if rowptr is not None: rowptr = torch.cat([rowptr, rowptr.new_full((diff_0, ), nnz)]) if rowcount is not None: rowcount = torch.cat([rowcount, rowcount.new_zeros(diff_0)]) elif diff_0 < 0: if rowptr is not None: rowptr = rowptr[:diff_0] if rowcount is not None: rowcount = rowcount[:diff_0] diff_1 = sparse_sizes[1] - old_sparse_sizes[1] colcount, colptr = self._colcount, self._colptr if diff_1 > 0: if colptr is not None: colptr = torch.cat([colptr, colptr.new_full((diff_1, ), nnz)]) if colcount is not None: colcount = torch.cat([colcount, colcount.new_zeros(diff_1)]) elif diff_1 < 0: if colptr is not None: colptr = colptr[:diff_1] if colcount is not None: colcount = colcount[:diff_1] return SparseStorage( row=self._row, rowptr=rowptr, col=self._col, value=self._value, sparse_sizes=sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=self._csr2csc, csc2csr=self._csc2csr, is_sorted=True, trust_data=True, ) def sparse_reshape(self, num_rows: int, num_cols: int): assert num_rows > 0 or num_rows == -1 assert num_cols > 0 or num_cols == -1 assert num_rows > 0 or num_cols > 0 total = self.sparse_size(0) * self.sparse_size(1) if num_rows == -1: num_rows = total // num_cols if num_cols == -1: num_cols = total // num_rows assert num_rows * num_cols == total idx = self.sparse_size(1) * self.row() + self.col() row = torch.div(idx, num_cols, rounding_mode='floor') col = idx % num_cols assert row.dtype == torch.long and col.dtype == torch.long return SparseStorage( row=row, rowptr=None, col=col, value=self._value, sparse_sizes=(num_rows, num_cols), rowcount=None, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True, trust_data=True, ) def has_rowcount(self) -> bool: return self._rowcount is not None def rowcount(self) -> torch.Tensor: rowcount = self._rowcount if rowcount is not None: return rowcount rowptr = self.rowptr() rowcount = rowptr[1:] - rowptr[:-1] self._rowcount = rowcount return rowcount def has_colptr(self) -> bool: return self._colptr is not None def colptr(self) -> torch.Tensor: colptr = self._colptr if colptr is not None: return colptr csr2csc = self._csr2csc if csr2csc is not None: colptr = torch.ops.torch_sparse.ind2ptr(self._col[csr2csc], self._sparse_sizes[1]) else: colptr = self._col.new_zeros(self._sparse_sizes[1] + 1) torch.cumsum(self.colcount(), dim=0, out=colptr[1:]) self._colptr = colptr return colptr def has_colcount(self) -> bool: return self._colcount is not None def colcount(self) -> torch.Tensor: colcount = self._colcount if colcount is not None: return colcount colptr = self._colptr if colptr is not None: colcount = colptr[1:] - colptr[:-1] else: colcount = scatter_add( torch.ones_like(self._col), self._col, dim_size=self._sparse_sizes[1], ) self._colcount = colcount return colcount def has_csr2csc(self) -> bool: return self._csr2csc is not None def csr2csc(self) -> torch.Tensor: csr2csc = self._csr2csc if csr2csc is not None: return csr2csc idx = self._sparse_sizes[0] * self._col + self.row() max_value = self._sparse_sizes[0] * self._sparse_sizes[1] _, csr2csc = index_sort(idx, max_value) self._csr2csc = csr2csc return csr2csc def has_csc2csr(self) -> bool: return self._csc2csr is not None def csc2csr(self) -> torch.Tensor: csc2csr = self._csc2csr if csc2csr is not None: return csc2csr max_value = self._sparse_sizes[0] * self._sparse_sizes[1] _, csc2csr = index_sort(self.csr2csc(), max_value) self._csc2csr = csc2csr return csc2csr def is_coalesced(self) -> bool: idx = self._col.new_full((self._col.numel() + 1, ), -1) idx[1:] = self._sparse_sizes[1] * self.row() + self._col return bool((idx[1:] > idx[:-1]).all()) def coalesce(self, reduce: str = "add"): idx = self._col.new_full((self._col.numel() + 1, ), -1) idx[1:] = self._sparse_sizes[1] * self.row() + self._col mask = idx[1:] > idx[:-1] if mask.all(): # Skip if indices are already coalesced. return self row = self.row()[mask] col = self._col[mask] value = self._value if value is not None: ptr = mask.nonzero().flatten() ptr = torch.cat([ptr, ptr.new_full((1, ), value.size(0))]) value = segment_csr(value, ptr, reduce=reduce) return SparseStorage( row=row, rowptr=None, col=col, value=value, sparse_sizes=self._sparse_sizes, rowcount=None, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True, trust_data=True, ) def fill_cache_(self): self.row() self.rowptr() self.rowcount() self.colptr() self.colcount() self.csr2csc() self.csc2csr() return self def clear_cache_(self): self._rowcount = None self._colptr = None self._colcount = None self._csr2csc = None self._csc2csr = None return self def cached_keys(self) -> List[str]: keys: List[str] = [] if self.has_rowcount(): keys.append('rowcount') if self.has_colptr(): keys.append('colptr') if self.has_colcount(): keys.append('colcount') if self.has_csr2csc(): keys.append('csr2csc') if self.has_csc2csr(): keys.append('csc2csr') return keys def num_cached_keys(self) -> int: return len(self.cached_keys()) def copy(self): return SparseStorage( row=self._row, rowptr=self._rowptr, col=self._col, value=self._value, sparse_sizes=self._sparse_sizes, rowcount=self._rowcount, colptr=self._colptr, colcount=self._colcount, csr2csc=self._csr2csc, csc2csr=self._csc2csr, is_sorted=True, trust_data=True, ) def clone(self): row = self._row if row is not None: row = row.clone() rowptr = self._rowptr if rowptr is not None: rowptr = rowptr.clone() col = self._col.clone() value = self._value if value is not None: value = value.clone() rowcount = self._rowcount if rowcount is not None: rowcount = rowcount.clone() colptr = self._colptr if colptr is not None: colptr = colptr.clone() colcount = self._colcount if colcount is not None: colcount = colcount.clone() csr2csc = self._csr2csc if csr2csc is not None: csr2csc = csr2csc.clone() csc2csr = self._csc2csr if csc2csr is not None: csc2csr = csc2csr.clone() return SparseStorage( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=self._sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True, trust_data=True, ) def type(self, dtype: torch.dtype, non_blocking: bool = False): value = self._value if value is not None: if dtype == value.dtype: return self else: return self.set_value( value.to(dtype=dtype, non_blocking=non_blocking), layout='coo', ) else: return self def type_as(self, tensor: torch.Tensor, non_blocking: bool = False): return self.type(dtype=tensor.dtype, non_blocking=non_blocking) def to_device(self, device: torch.device, non_blocking: bool = False): if device == self._col.device: return self row = self._row if row is not None: row = row.to(device, non_blocking=non_blocking) rowptr = self._rowptr if rowptr is not None: rowptr = rowptr.to(device, non_blocking=non_blocking) col = self._col.to(device, non_blocking=non_blocking) value = self._value if value is not None: value = value.to(device, non_blocking=non_blocking) rowcount = self._rowcount if rowcount is not None: rowcount = rowcount.to(device, non_blocking=non_blocking) colptr = self._colptr if colptr is not None: colptr = colptr.to(device, non_blocking=non_blocking) colcount = self._colcount if colcount is not None: colcount = colcount.to(device, non_blocking=non_blocking) csr2csc = self._csr2csc if csr2csc is not None: csr2csc = csr2csc.to(device, non_blocking=non_blocking) csc2csr = self._csc2csr if csc2csr is not None: csc2csr = csc2csr.to(device, non_blocking=non_blocking) return SparseStorage( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=self._sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True, trust_data=True, ) def device_as(self, tensor: torch.Tensor, non_blocking: bool = False): return self.to_device(device=tensor.device, non_blocking=non_blocking) def cuda(self): new_col = self._col.cuda() if new_col.device == self._col.device: return self row = self._row if row is not None: row = row.cuda() rowptr = self._rowptr if rowptr is not None: rowptr = rowptr.cuda() value = self._value if value is not None: value = value.cuda() rowcount = self._rowcount if rowcount is not None: rowcount = rowcount.cuda() colptr = self._colptr if colptr is not None: colptr = colptr.cuda() colcount = self._colcount if colcount is not None: colcount = colcount.cuda() csr2csc = self._csr2csc if csr2csc is not None: csr2csc = csr2csc.cuda() csc2csr = self._csc2csr if csc2csr is not None: csc2csr = csc2csr.cuda() return SparseStorage( row=row, rowptr=rowptr, col=new_col, value=value, sparse_sizes=self._sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True, trust_data=True, ) def pin_memory(self): row = self._row if row is not None: row = row.pin_memory() rowptr = self._rowptr if rowptr is not None: rowptr = rowptr.pin_memory() col = self._col.pin_memory() value = self._value if value is not None: value = value.pin_memory() rowcount = self._rowcount if rowcount is not None: rowcount = rowcount.pin_memory() colptr = self._colptr if colptr is not None: colptr = colptr.pin_memory() colcount = self._colcount if colcount is not None: colcount = colcount.pin_memory() csr2csc = self._csr2csc if csr2csc is not None: csr2csc = csr2csc.pin_memory() csc2csr = self._csc2csr if csc2csr is not None: csc2csr = csc2csr.pin_memory() return SparseStorage( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=self._sparse_sizes, rowcount=rowcount, colptr=colptr, colcount=colcount, csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True, trust_data=True, ) def is_pinned(self) -> bool: is_pinned = True row = self._row if row is not None: is_pinned = is_pinned and row.is_pinned() rowptr = self._rowptr if rowptr is not None: is_pinned = is_pinned and rowptr.is_pinned() is_pinned = self._col.is_pinned() value = self._value if value is not None: is_pinned = is_pinned and value.is_pinned() rowcount = self._rowcount if rowcount is not None: is_pinned = is_pinned and rowcount.is_pinned() colptr = self._colptr if colptr is not None: is_pinned = is_pinned and colptr.is_pinned() colcount = self._colcount if colcount is not None: is_pinned = is_pinned and colcount.is_pinned() csr2csc = self._csr2csc if csr2csc is not None: is_pinned = is_pinned and csr2csc.is_pinned() csc2csr = self._csc2csr if csc2csr is not None: is_pinned = is_pinned and csc2csr.is_pinned() return is_pinned def share_memory_(self) -> SparseStorage: row = self._row if row is not None: row.share_memory_() rowptr = self._rowptr if rowptr is not None: rowptr.share_memory_() self._col.share_memory_() value = self._value if value is not None: value.share_memory_() rowcount = self._rowcount if rowcount is not None: rowcount.share_memory_() colptr = self._colptr if colptr is not None: colptr.share_memory_() colcount = self._colcount if colcount is not None: colcount.share_memory_() csr2csc = self._csr2csc if csr2csc is not None: csr2csc.share_memory_() csc2csr = self._csc2csr if csc2csr is not None: csc2csr.share_memory_() def is_shared(self) -> bool: is_shared = True row = self._row if row is not None: is_shared = is_shared and row.is_shared() rowptr = self._rowptr if rowptr is not None: is_shared = is_shared and rowptr.is_shared() is_shared = is_shared and self._col.is_shared() value = self._value if value is not None: is_shared = is_shared and value.is_shared() rowcount = self._rowcount if rowcount is not None: is_shared = is_shared and rowcount.is_shared() colptr = self._colptr if colptr is not None: is_shared = is_shared and colptr.is_shared() colcount = self._colcount if colcount is not None: is_shared = is_shared and colcount.is_shared() csr2csc = self._csr2csc if csr2csc is not None: is_shared = is_shared and csr2csc.is_shared() csc2csr = self._csc2csr if csc2csr is not None: is_shared = is_shared and csc2csr.is_shared() return is_shared SparseStorage.share_memory_ = share_memory_ SparseStorage.is_shared = is_shared pytorch_sparse-0.6.18/torch_sparse/tensor.py000066400000000000000000000551231450761466100212530ustar00rootroot00000000000000from textwrap import indent from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import scipy.sparse import torch from torch_scatter import segment_csr from torch_sparse.storage import SparseStorage, get_layout @torch.jit.script class SparseTensor(object): storage: SparseStorage def __init__( self, row: Optional[torch.Tensor] = None, rowptr: Optional[torch.Tensor] = None, col: Optional[torch.Tensor] = None, value: Optional[torch.Tensor] = None, sparse_sizes: Optional[Tuple[Optional[int], Optional[int]]] = None, is_sorted: bool = False, trust_data: bool = False, ): self.storage = SparseStorage( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=None, colptr=None, colcount=None, csr2csc=None, csc2csr=None, is_sorted=is_sorted, trust_data=trust_data, ) @classmethod def from_storage(self, storage: SparseStorage): out = SparseTensor( row=storage._row, rowptr=storage._rowptr, col=storage._col, value=storage._value, sparse_sizes=storage._sparse_sizes, is_sorted=True, trust_data=True, ) out.storage._rowcount = storage._rowcount out.storage._colptr = storage._colptr out.storage._colcount = storage._colcount out.storage._csr2csc = storage._csr2csc out.storage._csc2csr = storage._csc2csr return out @classmethod def from_edge_index( self, edge_index: torch.Tensor, edge_attr: Optional[torch.Tensor] = None, sparse_sizes: Optional[Tuple[Optional[int], Optional[int]]] = None, is_sorted: bool = False, trust_data: bool = False, ): return SparseTensor( row=edge_index[0], rowptr=None, col=edge_index[1], value=edge_attr, sparse_sizes=sparse_sizes, is_sorted=is_sorted, trust_data=trust_data, ) @classmethod def from_dense(self, mat: torch.Tensor, has_value: bool = True): if mat.dim() > 2: index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero() else: index = mat.nonzero() index = index.t() row = index[0] col = index[1] value: Optional[torch.Tensor] = None if has_value: value = mat[row, col] return SparseTensor( row=row, rowptr=None, col=col, value=value, sparse_sizes=(mat.size(0), mat.size(1)), is_sorted=True, trust_data=True, ) @classmethod def from_torch_sparse_coo_tensor( self, mat: torch.Tensor, has_value: bool = True, ): mat = mat.coalesce() index = mat._indices() row, col = index[0], index[1] value: Optional[torch.Tensor] = None if has_value: value = mat.values() return SparseTensor( row=row, rowptr=None, col=col, value=value, sparse_sizes=(mat.size(0), mat.size(1)), is_sorted=True, trust_data=True, ) @classmethod def from_torch_sparse_csr_tensor( self, mat: torch.Tensor, has_value: bool = True, ): rowptr = mat.crow_indices() col = mat.col_indices() value: Optional[torch.Tensor] = None if has_value: value = mat.values() return SparseTensor( row=None, rowptr=rowptr, col=col, value=value, sparse_sizes=(mat.size(0), mat.size(1)), is_sorted=True, trust_data=True, ) @classmethod def eye(self, M: int, N: Optional[int] = None, has_value: bool = True, dtype: Optional[int] = None, device: Optional[torch.device] = None, fill_cache: bool = False): N = M if N is None else N row = torch.arange(min(M, N), device=device) col = row rowptr = torch.arange(M + 1, device=row.device) if M > N: rowptr[N + 1:] = N value: Optional[torch.Tensor] = None if has_value: value = torch.ones(row.numel(), dtype=dtype, device=row.device) rowcount: Optional[torch.Tensor] = None colptr: Optional[torch.Tensor] = None colcount: Optional[torch.Tensor] = None csr2csc: Optional[torch.Tensor] = None csc2csr: Optional[torch.Tensor] = None if fill_cache: rowcount = torch.ones(M, dtype=torch.long, device=row.device) if M > N: rowcount[N:] = 0 colptr = torch.arange(N + 1, dtype=torch.long, device=row.device) colcount = torch.ones(N, dtype=torch.long, device=row.device) if N > M: colptr[M + 1:] = M colcount[M:] = 0 csr2csc = csc2csr = row out = SparseTensor( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=(M, N), is_sorted=True, trust_data=True, ) out.storage._rowcount = rowcount out.storage._colptr = colptr out.storage._colcount = colcount out.storage._csr2csc = csr2csc out.storage._csc2csr = csc2csr return out def copy(self): return self.from_storage(self.storage) def clone(self): return self.from_storage(self.storage.clone()) def type(self, dtype: torch.dtype, non_blocking: bool = False): value = self.storage.value() if value is None or dtype == value.dtype: return self return self.from_storage( self.storage.type(dtype=dtype, non_blocking=non_blocking)) def type_as(self, tensor: torch.Tensor, non_blocking: bool = False): return self.type(dtype=tensor.dtype, non_blocking=non_blocking) def to_device(self, device: torch.device, non_blocking: bool = False): if device == self.device(): return self return self.from_storage( self.storage.to_device(device=device, non_blocking=non_blocking)) def device_as(self, tensor: torch.Tensor, non_blocking: bool = False): return self.to_device(device=tensor.device, non_blocking=non_blocking) # Formats ################################################################# def coo(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: return self.storage.row(), self.storage.col(), self.storage.value() def csr(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: return self.storage.rowptr(), self.storage.col(), self.storage.value() def csc(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: perm = self.storage.csr2csc() value = self.storage.value() if value is not None: value = value[perm] return self.storage.colptr(), self.storage.row()[perm], value # Storage inheritance ##################################################### def has_value(self) -> bool: return self.storage.has_value() def set_value_( self, value: Optional[torch.Tensor], layout: Optional[str] = None, ): self.storage.set_value_(value, layout) return self def set_value( self, value: Optional[torch.Tensor], layout: Optional[str] = None, ): return self.from_storage(self.storage.set_value(value, layout)) def sparse_sizes(self) -> Tuple[int, int]: return self.storage.sparse_sizes() def sparse_size(self, dim: int) -> int: return self.storage.sparse_sizes()[dim] def sparse_resize(self, sparse_sizes: Tuple[int, int]): return self.from_storage(self.storage.sparse_resize(sparse_sizes)) def sparse_reshape(self, num_rows: int, num_cols: int): return self.from_storage( self.storage.sparse_reshape(num_rows, num_cols)) def is_coalesced(self) -> bool: return self.storage.is_coalesced() def coalesce(self, reduce: str = "sum"): return self.from_storage(self.storage.coalesce(reduce)) def fill_cache_(self): self.storage.fill_cache_() return self def clear_cache_(self): self.storage.clear_cache_() return self def __eq__(self, other) -> bool: if not isinstance(other, self.__class__): return False if self.sizes() != other.sizes(): return False rowptrA, colA, valueA = self.csr() rowptrB, colB, valueB = other.csr() if valueA is None and valueB is not None: return False if valueA is not None and valueB is None: return False if not torch.equal(rowptrA, rowptrB): return False if not torch.equal(colA, colB): return False if valueA is None and valueB is None: return True return torch.equal(valueA, valueB) # Utility functions ####################################################### def fill_value_(self, fill_value: float, dtype: Optional[int] = None): value = torch.full( (self.nnz(), ), fill_value, dtype=dtype, device=self.device(), ) return self.set_value_(value, layout='coo') def fill_value(self, fill_value: float, dtype: Optional[int] = None): value = torch.full( (self.nnz(), ), fill_value, dtype=dtype, device=self.device(), ) return self.set_value(value, layout='coo') def sizes(self) -> List[int]: sparse_sizes = self.sparse_sizes() value = self.storage.value() if value is not None: return list(sparse_sizes) + list(value.size())[1:] else: return list(sparse_sizes) def size(self, dim: int) -> int: return self.sizes()[dim] def dim(self) -> int: return len(self.sizes()) def nnz(self) -> int: return self.storage.col().numel() def numel(self) -> int: value = self.storage.value() if value is not None: return value.numel() else: return self.nnz() def density(self) -> float: if self.sparse_size(0) == 0 or self.sparse_size(1) == 0: return 0.0 return self.nnz() / (self.sparse_size(0) * self.sparse_size(1)) def sparsity(self) -> float: return 1 - self.density() def avg_row_length(self) -> float: return self.nnz() / self.sparse_size(0) def avg_col_length(self) -> float: return self.nnz() / self.sparse_size(1) def bandwidth(self) -> int: row, col, _ = self.coo() return int((row - col).abs_().max()) def avg_bandwidth(self) -> float: row, col, _ = self.coo() return float((row - col).abs_().to(torch.float).mean()) def bandwidth_proportion(self, bandwidth: int) -> float: row, col, _ = self.coo() tmp = (row - col).abs_() return int((tmp <= bandwidth).sum()) / self.nnz() def is_quadratic(self) -> bool: return self.sparse_size(0) == self.sparse_size(1) def is_symmetric(self) -> bool: if not self.is_quadratic(): return False rowptr, col, value1 = self.csr() colptr, row, value2 = self.csc() if (rowptr != colptr).any() or (col != row).any(): return False if value1 is None or value2 is None: return True else: return bool((value1 == value2).all()) def to_symmetric(self, reduce: str = "sum"): N = max(self.size(0), self.size(1)) row, col, value = self.coo() idx = col.new_full((2 * col.numel() + 1, ), -1) idx[1:row.numel() + 1] = row idx[row.numel() + 1:] = col idx[1:] *= N idx[1:row.numel() + 1] += col idx[row.numel() + 1:] += row idx, perm = idx.sort() mask = idx[1:] > idx[:-1] perm = perm[1:].sub_(1) idx = perm[mask] if value is not None: ptr = mask.nonzero().flatten() ptr = torch.cat([ptr, ptr.new_full((1, ), perm.size(0))]) value = torch.cat([value, value])[perm] value = segment_csr(value, ptr, reduce=reduce) new_row = torch.cat([row, col], dim=0)[idx] new_col = torch.cat([col, row], dim=0)[idx] out = SparseTensor( row=new_row, rowptr=None, col=new_col, value=value, sparse_sizes=(N, N), is_sorted=True, trust_data=True, ) return out def detach_(self): value = self.storage.value() if value is not None: value.detach_() return self def detach(self): value = self.storage.value() if value is not None: value = value.detach() return self.set_value(value, layout='coo') def requires_grad(self) -> bool: value = self.storage.value() if value is not None: return value.requires_grad else: return False def requires_grad_( self, requires_grad: bool = True, dtype: Optional[int] = None, ): if requires_grad and not self.has_value(): self.fill_value_(1., dtype) value = self.storage.value() if value is not None: value.requires_grad_(requires_grad) return self def pin_memory(self): return self.from_storage(self.storage.pin_memory()) def is_pinned(self) -> bool: return self.storage.is_pinned() def device(self): return self.storage.col().device def cpu(self): return self.to_device(device=torch.device('cpu'), non_blocking=False) def cuda(self): return self.from_storage(self.storage.cuda()) def is_cuda(self) -> bool: return self.storage.col().is_cuda def dtype(self): value = self.storage.value() return value.dtype if value is not None else torch.float def is_floating_point(self) -> bool: value = self.storage.value() return torch.is_floating_point(value) if value is not None else True def bfloat16(self): return self.type(dtype=torch.bfloat16, non_blocking=False) def bool(self): return self.type(dtype=torch.bool, non_blocking=False) def byte(self): return self.type(dtype=torch.uint8, non_blocking=False) def char(self): return self.type(dtype=torch.int8, non_blocking=False) def half(self): return self.type(dtype=torch.half, non_blocking=False) def float(self): return self.type(dtype=torch.float, non_blocking=False) def double(self): return self.type(dtype=torch.double, non_blocking=False) def short(self): return self.type(dtype=torch.short, non_blocking=False) def int(self): return self.type(dtype=torch.int, non_blocking=False) def long(self): return self.type(dtype=torch.long, non_blocking=False) # Conversions ############################################################# def to_dense(self, dtype: Optional[int] = None) -> torch.Tensor: row, col, value = self.coo() if value is not None: mat = torch.zeros( self.sizes(), dtype=value.dtype, device=self.device(), ) else: mat = torch.zeros(self.sizes(), dtype=dtype, device=self.device()) if value is not None: mat[row, col] = value else: mat[row, col] = torch.ones( self.nnz(), dtype=mat.dtype, device=mat.device, ) return mat def to_torch_sparse_coo_tensor( self, dtype: Optional[int] = None, ) -> torch.Tensor: row, col, value = self.coo() index = torch.stack([row, col], dim=0) if value is None: value = torch.ones(self.nnz(), dtype=dtype, device=self.device()) return torch.sparse_coo_tensor(index, value, self.sizes()) def to_torch_sparse_csr_tensor( self, dtype: Optional[int] = None, ) -> torch.Tensor: rowptr, col, value = self.csr() if value is None: value = torch.ones(self.nnz(), dtype=dtype, device=self.device()) return torch.sparse_csr_tensor(rowptr, col, value, self.sizes()) def to_torch_sparse_csc_tensor( self, dtype: Optional[int] = None, ) -> torch.Tensor: colptr, row, value = self.csc() if value is None: value = torch.ones(self.nnz(), dtype=dtype, device=self.device()) return torch.sparse_csc_tensor(colptr, row, value, self.sizes()) # Python Bindings ############################################################# def share_memory_(self: SparseTensor) -> SparseTensor: self.storage.share_memory_() return self def is_shared(self: SparseTensor) -> bool: return self.storage.is_shared() def to(self, *args: Optional[List[Any]], **kwargs: Optional[Dict[str, Any]]) -> SparseTensor: device, dtype, non_blocking = torch._C._nn._parse_to(*args, **kwargs)[:3] if dtype is not None: self = self.type(dtype=dtype, non_blocking=non_blocking) if device is not None: self = self.to_device(device=device, non_blocking=non_blocking) return self def cpu(self) -> SparseTensor: return self.device_as(torch.tensor(0., device='cpu')) def cuda( self, device: Optional[Union[int, str]] = None, non_blocking: bool = False, ): return self.device_as(torch.tensor(0., device=device or 'cuda')) def __getitem__(self: SparseTensor, index: Any) -> SparseTensor: index = list(index) if isinstance(index, tuple) else [index] # More than one `Ellipsis` is not allowed... if len([ i for i in index if not isinstance(i, (torch.Tensor, np.ndarray)) and i == ... ]) > 1: raise SyntaxError dim = 0 out = self while len(index) > 0: item = index.pop(0) if isinstance(item, (list, tuple)): item = torch.tensor(item, device=self.device()) if isinstance(item, np.ndarray): item = torch.from_numpy(item).to(self.device()) if isinstance(item, int): out = out.select(dim, item) dim += 1 elif isinstance(item, slice): if item.step is not None: raise ValueError('Step parameter not yet supported.') start = 0 if item.start is None else item.start start = self.size(dim) + start if start < 0 else start stop = self.size(dim) if item.stop is None else item.stop stop = self.size(dim) + stop if stop < 0 else stop out = out.narrow(dim, start, max(stop - start, 0)) dim += 1 elif torch.is_tensor(item): if item.dtype == torch.bool: out = out.masked_select(dim, item) dim += 1 elif item.dtype == torch.long: out = out.index_select(dim, item) dim += 1 elif item == Ellipsis: if self.dim() - len(index) < dim: raise SyntaxError dim = self.dim() - len(index) else: raise SyntaxError return out def __repr__(self: SparseTensor) -> str: i = ' ' * 6 row, col, value = self.coo() infos = [] infos += [f'row={indent(row.__repr__(), i)[len(i):]}'] infos += [f'col={indent(col.__repr__(), i)[len(i):]}'] if value is not None: infos += [f'val={indent(value.__repr__(), i)[len(i):]}'] infos += [ f'size={tuple(self.sizes())}, nnz={self.nnz()}, ' f'density={100 * self.density():.02f}%' ] infos = ',\n'.join(infos) i = ' ' * (len(self.__class__.__name__) + 1) return f'{self.__class__.__name__}({indent(infos, i)[len(i):]})' SparseTensor.share_memory_ = share_memory_ SparseTensor.is_shared = is_shared SparseTensor.to = to SparseTensor.cpu = cpu SparseTensor.cuda = cuda SparseTensor.__getitem__ = __getitem__ SparseTensor.__repr__ = __repr__ # Scipy Conversions ########################################################### ScipySparseMatrix = Union[scipy.sparse.coo_matrix, scipy.sparse.csr_matrix, scipy.sparse.csc_matrix] @torch.jit.ignore def from_scipy(mat: ScipySparseMatrix, has_value: bool = True) -> SparseTensor: colptr = None if isinstance(mat, scipy.sparse.csc_matrix): colptr = torch.from_numpy(mat.indptr).to(torch.long) mat = mat.tocsr() rowptr = torch.from_numpy(mat.indptr).to(torch.long) mat = mat.tocoo() row = torch.from_numpy(mat.row).to(torch.long) col = torch.from_numpy(mat.col).to(torch.long) value = None if has_value: value = torch.from_numpy(mat.data) sparse_sizes = mat.shape[:2] storage = SparseStorage( row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=sparse_sizes, rowcount=None, colptr=colptr, colcount=None, csr2csc=None, csc2csr=None, is_sorted=True, ) return SparseTensor.from_storage(storage) @torch.jit.ignore def to_scipy( self: SparseTensor, layout: Optional[str] = None, dtype: Optional[torch.dtype] = None, ) -> ScipySparseMatrix: assert self.dim() == 2 layout = get_layout(layout) if not self.has_value(): ones = torch.ones(self.nnz(), dtype=dtype).numpy() if layout == 'coo': row, col, value = self.coo() row = row.detach().cpu().numpy() col = col.detach().cpu().numpy() value = value.detach().cpu().numpy() if self.has_value() else ones return scipy.sparse.coo_matrix((value, (row, col)), self.sizes()) elif layout == 'csr': rowptr, col, value = self.csr() rowptr = rowptr.detach().cpu().numpy() col = col.detach().cpu().numpy() value = value.detach().cpu().numpy() if self.has_value() else ones return scipy.sparse.csr_matrix((value, col, rowptr), self.sizes()) elif layout == 'csc': colptr, row, value = self.csc() colptr = colptr.detach().cpu().numpy() row = row.detach().cpu().numpy() value = value.detach().cpu().numpy() if self.has_value() else ones return scipy.sparse.csc_matrix((value, row, colptr), self.sizes()) SparseTensor.from_scipy = from_scipy SparseTensor.to_scipy = to_scipy pytorch_sparse-0.6.18/torch_sparse/testing.py000066400000000000000000000013321450761466100214070ustar00rootroot00000000000000from typing import Any import torch import torch_scatter from packaging import version reductions = ['sum', 'add', 'mean', 'min', 'max'] dtypes = [torch.half, torch.float, torch.double, torch.int, torch.long] grad_dtypes = [torch.half, torch.float, torch.double] if version.parse(torch_scatter.__version__) > version.parse("2.0.9"): dtypes.append(torch.bfloat16) grad_dtypes.append(torch.bfloat16) devices = [torch.device('cpu')] if torch.cuda.is_available(): devices += [torch.device('cuda:0')] if torch.backends.mps.is_available(): devices += [torch.device('mps')] def tensor(x: Any, dtype: torch.dtype, device: torch.device): return None if x is None else torch.tensor(x, dtype=dtype, device=device) pytorch_sparse-0.6.18/torch_sparse/transpose.py000066400000000000000000000034551450761466100217600ustar00rootroot00000000000000import torch from torch_sparse.storage import SparseStorage from torch_sparse.tensor import SparseTensor def t(src: SparseTensor) -> SparseTensor: csr2csc = src.storage.csr2csc() row, col, value = src.coo() if value is not None: value = value[csr2csc] sparse_sizes = src.storage.sparse_sizes() storage = SparseStorage( row=col[csr2csc], rowptr=src.storage._colptr, col=row[csr2csc], value=value, sparse_sizes=(sparse_sizes[1], sparse_sizes[0]), rowcount=src.storage._colcount, colptr=src.storage._rowptr, colcount=src.storage._rowcount, csr2csc=src.storage._csc2csr, csc2csr=csr2csc, is_sorted=True, ) return src.from_storage(storage) SparseTensor.t = lambda self: t(self) ############################################################################### def transpose(index, value, m, n, coalesced=True): """Transposes dimensions 0 and 1 of a sparse tensor. Args: index (:class:`LongTensor`): The index tensor of sparse matrix. value (:class:`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 :obj:`False`, will not coalesce the output. (default: :obj:`True`) :rtype: (:class:`LongTensor`, :class:`Tensor`) """ row, col = index row, col = col, row if coalesced: sparse_sizes = (n, m) storage = SparseStorage(row=row, col=col, value=value, sparse_sizes=sparse_sizes, is_sorted=False) storage = storage.coalesce() row, col, value = storage.row(), storage.col(), storage.value() return torch.stack([row, col], dim=0), value pytorch_sparse-0.6.18/torch_sparse/typing.py000066400000000000000000000003171450761466100212460ustar00rootroot00000000000000try: import pyg_lib # noqa WITH_PYG_LIB = True WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort') except ImportError: pyg_lib = object WITH_PYG_LIB = False WITH_INDEX_SORT = False pytorch_sparse-0.6.18/torch_sparse/utils.py000066400000000000000000000013311450761466100210710ustar00rootroot00000000000000from typing import Any, Optional, Tuple import torch import torch_sparse.typing from torch_sparse.typing import pyg_lib try: from typing_extensions import Final # noqa except ImportError: from torch.jit import Final # noqa def index_sort( inputs: torch.Tensor, max_value: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: r"""See pyg-lib documentation for more details: https://pyg-lib.readthedocs.io/en/latest/modules/ops.html""" if not torch_sparse.typing.WITH_INDEX_SORT: # pragma: no cover return inputs.sort() return pyg_lib.ops.index_sort(inputs, max_value) def is_scalar(other: Any) -> bool: return isinstance(other, int) or isinstance(other, float)