Merge branch 'docker_january_2025' into checkhealth

Author: jeanfeydy

.cirrus.yml

@@ -13,6 +13,7 @@ task:
   name: Tests (macOS)
   macos_instance:
     image: ghcr.io/cirruslabs/macos-ventura-xcode:latest
+
   script:
     - brew install libomp
     - brew install [email protected]
@@ -21,3 +22,4 @@ task:
     - pip install --upgrade pip
     - sh ./pytest.sh
 
+

.github/workflows/cuda_test.yml

@@ -0,0 +1,15 @@
+# Launch test on a cuda machine
+name: Cuda test
+
+on:
+  push
+
+jobs:
+  test:
+    runs-on: self-hosted
+    steps:
+      - name: Clone Repository
+        uses: actions/checkout@v2
+      - name: Test with pytest
+        run: |
+          sh ./pytest.sh

Dockerfile

@@ -15,6 +15,8 @@
 # --tag getkeops/keops:2.1-geomloss0.2.5-cuda11.3-pytorch1.11-full \
 # --no-cache .
 
+ARG NUMPY_VERSION=1.26.4
+
 # KeOps version - the most important parameter:
 ARG KEOPS_VERSION=2.2
 # We also include all the libraries hosted on www.kernel-operations.io,
@@ -26,7 +28,7 @@ ARG GEOMLOSS_VERSION=0.2.5
 
 
 # Base OS:
-ARG BASE_IMAGE=ubuntu:22.04
+ARG BASE_IMAGE=ubuntu:24.04
 # Useful to test support across Python versions:
 ARG PYTHON_VERSION=3.10
 
@@ -113,7 +115,7 @@ RUN curl -fsSL -v -o ~/miniconda.sh -O  https://repo.anaconda.com/miniconda/Mini
     /opt/conda/bin/conda install -y python=${PYTHON_VERSION} \
         conda-build \
         pyyaml \
-        numpy \
+        numpy=${NUMPY_VERSION} \
         ipython \
         matplotlib \
         ipykernel && \
@@ -188,11 +190,13 @@ RUN /opt/conda/bin/pip install -r doc-requirements.txt
 
 
 # Super-full environment with optional dependencies:
-FROM keops-doc as keops-full
+FROM keops-doc AS keops-full
 # PyTorch-scatter is a complex dependency:
 # it relies on binaries that often lag behind new PyTorch releases
 # by a few days/weeks.
 ARG PYTORCH_SCATTER_VERSION
-RUN /opt/conda/bin/conda install -y -c pyg \
-    pytorch-scatter==${PYTORCH_SCATTER_VERSION} && \
-    /opt/conda/bin/conda clean -ya
+RUN /opt/conda/bin/pip install torch-scatter -f ${PYTORCH_SCATTER_VERSION}
+
+#RUN /opt/conda/bin/conda install -y -c pyg \
+#    pytorch-scatter==${PYTORCH_SCATTER_VERSION} && \
+#    /opt/conda/bin/conda clean -ya

docker-images.sh

@@ -5,21 +5,22 @@
 # to clear the docker cache and make sure that the config works
 # with the current version of e.g. the Ubuntu repositories.
 
-# Up to date as of Monday, Feb. 12th, 2024:
-PYTHON_VERSION=3.11
-KEOPS_VERSION=2.2.2
+# Up to date as of Monday, Jan. 7th, 2025:
+PYTHON_VERSION=3.12
+NUMPY_VERSION=2.2.1
+KEOPS_VERSION=2.2.3
 GEOMLOSS_VERSION=0.2.6
-CUDA_VERSION=12.1
-CUDA_CHANNEL=nvidia/label/cuda-12.1.1
-PYTORCH_VERSION=2.2.0
-TORCHVISION_VERSION=0.17.0
-TORCHAUDIO_VERSION=2.2.0
-PYTORCH_SCATTER_VERSION=2.1.2
-PYTEST_VERSION=8.0.0
-HYPOTHESIS_VERSION=6.98.4
-JAXTYPING_VERSION=0.2.25
-BEARTYPE_VERSION=0.17.1
-BLACK_VERSION=24.1.1
+CUDA_VERSION=12.4
+CUDA_CHANNEL=nvidia/label/cuda-12.4.1
+PYTORCH_VERSION=2.5.1
+TORCHVISION_VERSION=0.20.1
+TORCHAUDIO_VERSION=2.5.1
+PYTORCH_SCATTER_VERSION=https://data.pyg.org/whl/torch-2.5.1+cu124.html
+PYTEST_VERSION=8.3.4
+HYPOTHESIS_VERSION=6.123.7
+JAXTYPING_VERSION=0.2.36
+BEARTYPE_VERSION=0.19.0
+BLACK_VERSION=24.10.0
 
 VERSION_TAG=${KEOPS_VERSION}-geomloss${GEOMLOSS_VERSION}-cuda${CUDA_VERSION}-pytorch${PYTORCH_VERSION}-python${PYTHON_VERSION}
 
@@ -28,6 +29,7 @@ do
     docker build \
     --target ${TARGET} \
     --build-arg PYTHON_VERSION=${PYTHON_VERSION} \
+    --build-arg NUMPY_VERSION=${NUMPY_VERSION} \
     --build-arg KEOPS_VERSION=${KEOPS_VERSION} \
     --build-arg GEOMLOSS_VERSION=${GEOMLOSS_VERSION} \
     --build-arg CUDA_VERSION=${CUDA_VERSION} \

keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp

@@ -635,13 +635,18 @@ template <typename TYPE> class KeOps_module {
 
     if (RR.tagRanges == 1) {
       CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)lookup_d));
+      // slices_x_d and ranges_y_d are always allocated in range_preprocess_from_host 
+      // if RR.tagRanges == 1, regardless of nbatchdims. So we always free them
+      CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)slices_x_d));
+      CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)ranges_y_d));
+
+      // offsets_d is only allocated if nbatchdims > 0, so free it with condition
       if (SS.nbatchdims > 0) {
-        CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)slices_x_d));
-        CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)ranges_y_d));
-        CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)offsets_d));
+          CUDA_SAFE_CALL(cuMemFree((CUdeviceptr)offsets_d));
       }
     }
 
+
     // end_ = end = clock();
     ////std::cout << "  time for last part : " << double(//end_ - start_) /
     /// CLOCKS_PER_SEC << std::endl; /std::cout << "time for launch_keops inner

pykeops/pykeops/benchmarks/benchmark_utils.py

@@ -340,15 +340,15 @@ def bench_config(
 
     except (TimeoutError, IndexError):  # Thrown by Nloops.pop(0) if Nloops = []
         print("** Too slow!")
-        not_recorded_times = (len(problem_sizes) - len(times)) * [np.Infinity]
+        not_recorded_times = (len(problem_sizes) - len(times)) * [np.inf]
 
     except NotImplementedError:
         print("** This metric is not supported!")
-        not_recorded_times = (len(problem_sizes) - len(times)) * [np.Infinity]
+        not_recorded_times = (len(problem_sizes) - len(times)) * [np.inf]
 
     except ValueError as err:
         print(err)
-        not_recorded_times = (len(problem_sizes) - len(times)) * [np.NINF]
+        not_recorded_times = (len(problem_sizes) - len(times)) * [-np.inf]
 
     except RuntimeError as err:
         print(err)

pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py

@@ -37,7 +37,6 @@
 
 from scipy.sparse import diags
 from scipy.sparse.linalg import aslinearoperator, cg
-from scipy.sparse.linalg.interface import IdentityOperator
 
 from pykeops.numpy import KernelSolve as KernelSolve_np, LazyTensor
 from pykeops.torch import KernelSolve

pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py

@@ -67,7 +67,6 @@
 
 from scipy.sparse import diags
 from scipy.sparse.linalg import aslinearoperator, cg
-from scipy.sparse.linalg.interface import IdentityOperator
 
 print("Solving a Gaussian linear system, with {} points in dimension {}.".format(N, D))
 start = time.time()

pykeops/pykeops/tutorials/backends/plot_scipy.py

@@ -154,10 +154,8 @@
 #       L_{\text{norm}}= \text{Id} - \text{diag}(D^{-1/2}) K \text{diag}(D^{-1/2}).
 
 
-from scipy.sparse.linalg.interface import IdentityOperator
-
 D_2 = aslinearoperator(diags(1 / np.sqrt(D)))
-L_norm = IdentityOperator((N, N)) - D_2 @ K @ D_2
+L_norm = aslinearoperator(diags(np.ones_like(D))) - D_2 @ K @ D_2
 L_norm.dtype = np.dtype(
     dtype
 )  # Scipy Bugfix: by default, "-" removes the dtype information...
@@ -373,7 +371,7 @@
 
 D = K @ np.ones(N, dtype=dtype)  # Sum along the lines of the adjacency matrix
 D_2 = aslinearoperator(diags(1 / np.sqrt(D)))
-L_norm = IdentityOperator((N, N)) - D_2 @ K @ D_2
+L_norm = aslinearoperator(diags(np.ones_like(D))) - D_2 @ K @ D_2
 L_norm.dtype = np.dtype(
     dtype
 )  # Scipy Bugfix: by default, "-" removes the dtype information...

readme.md

@@ -5,8 +5,8 @@
 <!-- [![Build status](https://ci.inria.fr/keops/buildStatus/icon?job=keops%2Fmain)](https://ci.inria.fr/keops/job/keops/job/main/) -->
 [![PyPI version](https://img.shields.io/pypi/v/pykeops?color=blue)](https://pypi.org/project/pykeops/)
 [![PyPI downloads](https://pepy.tech/badge/pykeops?color=green)](https://pypi.org/project/pykeops/)
-[![CRAN version](https://img.shields.io/cran/v/rkeops?color=yellowgreen)](https://cran.r-project.org/web/packages/rkeops/index.html)
-[![CRAN downloads](http://cranlogs.r-pkg.org/badges/grand-total/rkeops?color=yellowgreen)](https://cran.r-project.org/web/packages/rkeops/index.html)
+<!-- [![CRAN version](https://img.shields.io/cran/v/rkeops?color=yellowgreen)](https://cran.r-project.org/web/packages/rkeops/index.html) -->
+<!-- [![CRAN downloads](http://cranlogs.r-pkg.org/badges/grand-total/rkeops?color=yellowgreen)](https://cran.r-project.org/web/packages/rkeops/index.html) -->
 
 
 Please visit our [website](https://www.kernel-operations.io/) for