CIP-0001 Phase 3: DeepGP psi_chunk_size / psi_threads

.github/workflows/tests.yml

@@ -33,7 +33,9 @@ jobs:
           python -m pip install --upgrade pip setuptools wheel
           # Install runtime + test deps explicitly: editable metadata from
           # setup.py is not reliably applied by current pip/setuptools.
-          python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables "coverage>=7.0"
+          python -m pip install "numpy>=1.7" "scipy>=0.12" matplotlib h5py tables "coverage>=7.0"
+          # CIP-0001 Phase 3 needs GPy psicomp chunk_size (merged in GPy #1299; not yet on PyPI).
+          python -m pip install "git+https://github.com/SheffieldML/GPy.git@main"
           python -m pip install -e . --no-deps
 
       - name: Run unit tests with coverage
@@ -63,7 +65,8 @@ jobs:
           sudo apt-get update
           sudo apt-get install -y --no-install-recommends libopenmpi-dev openmpi-bin
           python -m pip install --upgrade pip setuptools wheel
-          python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables mpi4py
+          python -m pip install "numpy>=1.7" "scipy>=0.12" matplotlib h5py tables mpi4py
+          python -m pip install "git+https://github.com/SheffieldML/GPy.git@main"
           python -m pip install -e . --no-deps
 
       - name: MPI parity (2 ranks vs serial)

README.md

@@ -21,6 +21,8 @@ To run the suite locally:
 python -m pip install -e ".[test]"
 # or, if editable extras are unreliable:
 python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables "coverage>=7.0"
+# For psi_chunk_size / psi_threads (CIP-0001 Phase 3) until the next GPy release:
+# python -m pip install "git+https://github.com/SheffieldML/GPy.git@main"
 python -m pip install -e . --no-deps
 
 python -m coverage run -m unittest discover -s deepgp/testing -p 'test*.py' -v
@@ -47,3 +49,15 @@ single-machine role of MPI. ``mpi_comm`` with ``back_constraint=False`` raises
 ``ValueError`` (that configuration produced silently wrong results).
 
 An optional GitHub Actions job runs a 2-rank parity check while MPI remains.
+
+Psi chunking (optional)
+-----------------------
+
+Large-N models can bound memory and use multiple cores via GPy's chunked RBF /
+linear psi-statistics (CIP-0001 Phase 3)::
+
+    m = deepgp.DeepGP(..., psi_chunk_size=500, psi_threads=4)
+
+Defaults (``psi_chunk_size=None``, ``psi_threads=1``) preserve previous
+behaviour. Threading runs inside GPy's pure psi helpers; PyDeepGP itself stays
+serial at the kernel API.

cip/cip0001.md

@@ -27,7 +27,7 @@ title: Replace Theano and MPI with modern NumPy
 
 - [x] Proposed - Initial idea documented (open for Discussion feedback)
 - [x] Accepted - Phase 1 started on branch `cip0001-numpy-encoder` (regularisation: keep no-op + FutureWarning)
-- [x] In Progress - Phase 2 (MPI correctness / deprecation) under way
+- [x] In Progress - Phase 3 (chunked psi-statistics) under way
 - [ ] Implemented - Work complete, awaiting verification
 - [ ] Closed - Verified and complete
 - [ ] Rejected - Will not be implemented
@@ -256,8 +256,8 @@ pluggable reducer in R14, so the MPI-specific code shrinks to that one place.
 - [ ] Discussion opened
 - [x] Accept (partial): regularisation kept as no-op with `FutureWarning` (R2)
 - [x] Phase 1: NumPy encoder (R1–R7) — merged (PR #37)
-- [x] Phase 2: MPI correctness and deprecation (R8–R11) on `cip0001-mpi-deprecate` — typemap removed; unsupported configs raise; FutureWarning; optional MPI CI job
-- [ ] Phase 3: chunked psi-statistics (R12 in GPy, R13–R16)
+- [x] Phase 2: MPI correctness and deprecation (R8–R11) — merged (PR #38)
+- [x] Phase 3: chunked psi-statistics — GPy R12/R13 merged (GPy PR #1299); PyDeepGP R14–R16 on `cip0001-psi-chunking` (`psi_chunk_size` / `psi_threads`)
 - [ ] Many-core benchmarks recorded
 - [ ] Phase 4: MPI removed (R17–R24) or reduced to the reducer hook
 

deepgp/inference/svi_vardtc.py

@@ -26,9 +26,11 @@ class SVI_VarDTC(LatentFunctionInference):
 
     """
     const_jitter = 1e-6
-    def __init__(self, mpi_comm=None, mpi_root=0):
+    def __init__(self, mpi_comm=None, mpi_root=0, psi_chunk_size=None, psi_threads=1):
         self.mpi_comm = mpi_comm
         self.mpi_root=mpi_root
+        self.psi_chunk_size = psi_chunk_size
+        self.psi_threads = 1 if psi_threads is None else int(psi_threads)
 
     def get_trYYT(self, Y):
         return np.sum(np.square(Y))
@@ -41,10 +43,34 @@ def get_YYTfactor(self, Y):
             return jitchol(tdot(Y))
 
     def gatherPsiStat(self, kern, X, Z, Y, beta, uncertain_inputs):
+        from ..util.psi_chunking import configure_kern_psicomp, row_slices
 
         assert beta.size == 1
 
-        if uncertain_inputs:
+        psicomp_chunked = configure_kern_psicomp(kern, self.psi_chunk_size, self.psi_threads)
+        N = X.shape[0]
+        chunk = self.psi_chunk_size
+        use_serial_chunks = (
+            uncertain_inputs
+            and chunk is not None
+            and int(chunk) > 0
+            and int(chunk) < N
+            and not psicomp_chunked
+        )
+
+        if use_serial_chunks:
+            psi1_parts = []
+            psi2 = None
+            psi0 = 0.0
+            for sl in row_slices(N, chunk):
+                Xc = X[sl]
+                psi0 = psi0 + kern.psi0(Z, Xc).sum()
+                psi1_parts.append(kern.psi1(Z, Xc) * beta)
+                psi2_c = kern.psi2(Z, Xc) * beta
+                psi2 = psi2_c if psi2 is None else psi2 + psi2_c
+            psi1 = np.vstack(psi1_parts)
+            psi0 = psi0 * beta
+        elif uncertain_inputs:
             psi0 = kern.psi0(Z, X).sum()*beta
             psi1 = kern.psi1(Z, X)*beta
             psi2 = kern.psi2(Z, X)*beta

deepgp/inference/vardtc.py

@@ -29,9 +29,11 @@ class VarDTC(LatentFunctionInference):
 
     """
     const_jitter = 1e-6
-    def __init__(self, mpi_comm=None):
+    def __init__(self, mpi_comm=None, psi_chunk_size=None, psi_threads=1):
 
         self.mpi_comm = mpi_comm
+        self.psi_chunk_size = psi_chunk_size
+        self.psi_threads = 1 if psi_threads is None else int(psi_threads)
         self.Y_speedup = False # Replace Y with the cholesky factor of YY.T, but the computation of posterior object will be skipped.
 
     def get_trYYT(self, Y):
@@ -45,20 +47,46 @@ def get_YYTfactor(self, Y):
             return jitchol(tdot(Y))
 
     def gatherPsiStat(self, kern, X, Z, Y, beta, uncertain_inputs):
+        from ..util.psi_chunking import configure_kern_psicomp, row_slices
 
         assert beta.size == 1
 
-        if uncertain_inputs:
+        # Prefer GPy R12 chunking/threading when the kernel supports it.
+        psicomp_chunked = configure_kern_psicomp(kern, self.psi_chunk_size, self.psi_threads)
+        N = X.shape[0]
+        chunk = self.psi_chunk_size
+        use_serial_chunks = (
+            uncertain_inputs
+            and chunk is not None
+            and int(chunk) > 0
+            and int(chunk) < N
+            and not psicomp_chunked
+        )
+
+        if use_serial_chunks:
+            # CIP-0001 R14: serial map-reduce over row chunks for any kernel.
+            psi1_parts = []
+            psi2 = None
+            psi0 = 0.0
+            for sl in row_slices(N, chunk):
+                Xc = X[sl]
+                psi0 = psi0 + kern.psi0(Z, Xc).sum()
+                psi1_parts.append(kern.psi1(Z, Xc) * beta)
+                psi2_c = kern.psi2(Z, Xc) * beta
+                psi2 = psi2_c if psi2 is None else psi2 + psi2_c
+            psi1 = np.vstack(psi1_parts)
+            psi0 = psi0 * beta
+        elif uncertain_inputs:
             psi0 = kern.psi0(Z, X)
             psi1 = kern.psi1(Z, X)*beta
             psi2 = kern.psi2(Z, X)*beta
+            psi0 = psi0.sum()*beta
         else:
             psi0 = kern.Kdiag(X)
             psi1 = kern.K(X, Z)
             psi2 = np.dot(psi1.T,psi1)*beta
             psi1 = psi1*beta
-
-        psi0 = psi0.sum()*beta
+            psi0 = psi0.sum()*beta
 
         if isinstance(Y, VariationalPosterior):
             m, s = Y.mean, Y.variance

deepgp/layers/layers.py

@@ -21,9 +21,13 @@ def __init__(self, X, Y, Z, kernel, likelihood,
                  normalizer=False,
                  mpi_comm=None,
                  mpi_root=0,
-                 auto_update=True):
+                 auto_update=True,
+                 psi_chunk_size=None,
+                 psi_threads=1):
         self.mpi_comm = mpi_comm
         self.mpi_root = mpi_root
+        self.psi_chunk_size = psi_chunk_size
+        self.psi_threads = 1 if psi_threads is None else int(psi_threads)
         self.psicov = False
         self.svi = False
         self.qU_ratio = 1.
@@ -32,7 +36,7 @@ def __init__(self, X, Y, Z, kernel, likelihood,
         if inference_method is None:
             from ..inference import VarDTC_parallel, VarDTC
             if mpi_comm is None:
-                inference_method = VarDTC()
+                inference_method = VarDTC(psi_chunk_size=psi_chunk_size, psi_threads=self.psi_threads)
             else:
                 inference_method = VarDTC_parallel(mpi_comm, mpi_root)
         elif inference_method=='inferentia' and mpi_comm is None:
@@ -41,7 +45,7 @@ def __init__(self, X, Y, Z, kernel, likelihood,
             self.psicov = True
         elif inference_method=='svi':
             from ..inference import SVI_VarDTC
-            inference_method = SVI_VarDTC()
+            inference_method = SVI_VarDTC(psi_chunk_size=psi_chunk_size, psi_threads=self.psi_threads)
             self.svi = True
 
         super(SparseGP_MPI, self).__init__(X, Y, Z, kernel,
@@ -50,6 +54,9 @@ def __init__(self, X, Y, Z, kernel, likelihood,
                                            inference_method=inference_method,
                                            name=name, Y_metadata=Y_metadata,
                                            normalizer=normalizer)
+
+        from ..util.psi_chunking import apply_psi_chunking
+        apply_psi_chunking(self, psi_chunk_size, self.psi_threads)
 
         if self.svi:
             from ..util.misc import comp_mapping
@@ -172,7 +179,8 @@ def __init__(self, layer_lower,
                  Z=None, num_inducing=10,  kernel=None,
                  inference_method=None, uncertain_inputs=True,
                  mpi_comm=None, mpi_root=0, back_constraint=True,
-                 encoder=None, auto_update=True, name='layer'):
+                 encoder=None, auto_update=True, name='layer',
+                 psi_chunk_size=None, psi_threads=1):
 
         self.uncertain_inputs = uncertain_inputs
         self.layer_lower = layer_lower
@@ -205,7 +213,9 @@ def __init__(self, layer_lower,
         super(Layer, self).__init__(X, Y, Z, kernel,
                                     likelihood, inference_method=inference_method,
                                     mpi_comm=mpi_comm, mpi_root=mpi_root,
-                                    auto_update=auto_update, name=name)
+                                    auto_update=auto_update, name=name,
+                                    psi_chunk_size=psi_chunk_size,
+                                    psi_threads=psi_threads)
         self.link_parameter(self.Kuu_sigma)
         if back_constraint: self.encoder = encoder
 
@@ -304,7 +314,8 @@ def __init__(self, dim_down, dim_up,
                     mpi_comm=None, mpi_root=0,
                  back_constraint=True, encoder=None,
                  auto_update=True, repeatX=False,
-                 repeatXsplit=0, name='obslayer'):
+                 repeatXsplit=0, name='obslayer',
+                 psi_chunk_size=None, psi_threads=1):
         self.dim_up, self.dim_down = dim_up, dim_down
         self._Y = Y
         self.repeatX = repeatX
@@ -321,7 +332,9 @@ def __init__(self, dim_down, dim_up,
                                             mpi_comm=mpi_comm, mpi_root=mpi_root,
                                             back_constraint=back_constraint,
                                             encoder=encoder, auto_update=auto_update,
-                                            name=name)
+                                            name=name,
+                                            psi_chunk_size=psi_chunk_size,
+                                            psi_threads=psi_threads)
 
     def set_as_toplayer(self, flag=True):
         if flag:
@@ -339,7 +352,8 @@ def __init__(self, layer_lower, dim_up,
                  kernel=None, inference_method=None,
                  noise_var=1e-2, init='rand',
                  mpi_comm=None, mpi_root=0, back_constraint=True,
-                 encoder=None, auto_update=True, name='hiddenlayer'):
+                 encoder=None, auto_update=True, name='hiddenlayer',
+                 psi_chunk_size=None, psi_threads=1):
 
         self.dim_up, self.dim_down = dim_up, layer_lower.X.shape[1] #self.Y.shape[1]
         likelihood = likelihoods.Gaussian(variance=noise_var)
@@ -354,7 +368,9 @@ def __init__(self, layer_lower, dim_up,
                                           mpi_comm=mpi_comm, mpi_root=mpi_root,
                                           back_constraint=back_constraint,
                                           encoder=encoder, auto_update=auto_update,
-                                          name=name)
+                                          name=name,
+                                          psi_chunk_size=psi_chunk_size,
+                                          psi_threads=psi_threads)
 
     def update_layer(self):
         super(HiddenLayer,self).update_layer()
@@ -398,7 +414,8 @@ def __init__(self, layer_lower, dim_up, X=None, X_variance=None, Z=None,
                  mpi_comm=None, mpi_root=0,
                  encoder=None,
                  back_constraint=True,
-                 auto_update=True, name='tophiddenlayer'):
+                 auto_update=True, name='tophiddenlayer',
+                 psi_chunk_size=None, psi_threads=1):
 
         self.dim_up, self.dim_down = dim_up, layer_lower.X.shape[1]
         likelihood = likelihoods.Gaussian(variance=noise_var)
@@ -411,7 +428,9 @@ def __init__(self, layer_lower, dim_up, X=None, X_variance=None, Z=None,
                                              mpi_comm=mpi_comm, mpi_root=mpi_root,
                                              back_constraint=back_constraint,
                                              encoder=encoder, auto_update=auto_update,
-                                             name=name)
+                                             name=name,
+                                             psi_chunk_size=psi_chunk_size,
+                                             psi_threads=psi_threads)
 
     def update_layer(self):
         super(TopHiddenLayer,self).update_layer()

deepgp/models/model.py

@@ -29,7 +29,8 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                  likelihood = None, inits='PCA', name='deepgp',
                  kernels=None, obs_data='cont', back_constraint=True,
                  encoder_dims=None, mpi_comm=None, mpi_root=0, repeatX=False,
-                 inference_method=None, **kwargs):
+                 inference_method=None, psi_chunk_size=None, psi_threads=1,
+                 **kwargs):
         super(DeepGP, self).__init__(name=name)
         self._IN_OPTIMIZATION_ = False
 
@@ -49,6 +50,8 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                 )
 
         self.mpi_comm, self.mpi_root = mpi_comm, mpi_root
+        self.psi_chunk_size = psi_chunk_size
+        self.psi_threads = 1 if psi_threads is None else int(psi_threads)
         self.obs_data = obs_data
         self.nLayers = len(nDims)-1
         self.nDims = nDims
@@ -128,7 +131,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                                                          mpi_root=mpi_root,
                                                          auto_update=auto_update,
                                                          repeatX=True,
-                                                         repeatXsplit=self.nDimsOrig[1]))
+                                                         repeatXsplit=self.nDimsOrig[1],
+                                                         psi_chunk_size=psi_chunk_size,
+                                                         psi_threads=self.psi_threads))
                         self.layers[-1].X_dim_free = range(self.nDimsOrig[1])
                         self.layers[-1].X_dim_top  = np.arange(self.nDimsOrig[1],
                                                                self.nDimsOrig[1]+X.shape[1]).tolist()
@@ -141,7 +146,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                                                          back_constraint=back_constraint,
                                                          inference_method=inference_method,
                                                          mpi_comm=mpi_comm, mpi_root=mpi_root,
-                                                         auto_update=auto_update))
+                                                         auto_update=auto_update,
+                                                         psi_chunk_size=psi_chunk_size,
+                                                         psi_threads=self.psi_threads))
             elif i==self.nLayers-1:
                 if isinstance(X, VariationalPosterior):
                     X_variance, uncertain_inputs = X.variance.values, True
@@ -158,7 +165,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                                                   back_constraint=True if uncertain_inputs and back_constraint else False,
                                                   inference_method=inference_method,
                                                   mpi_comm=mpi_comm, mpi_root=mpi_root,
-                                                  auto_update=auto_update))
+                                                  auto_update=auto_update,
+                                                  psi_chunk_size=psi_chunk_size,
+                                                  psi_threads=self.psi_threads))
             else:
                 self.layers.append(HiddenLayer(self.layers[-1],
                                                nDims[i+1], X=Xs[i],
@@ -168,7 +177,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10,
                                                back_constraint=back_constraint,
                                                inference_method=inference_method,
                                                mpi_comm=mpi_comm, mpi_root=mpi_root,
-                                               auto_update=auto_update)) 
+                                               auto_update=auto_update,
+                                               psi_chunk_size=psi_chunk_size,
+                                               psi_threads=self.psi_threads))
         if self.nLayers==1:
             self.layers[0].set_as_toplayer()
 
@@ -294,7 +305,9 @@ def add_layer(self, dim_up, X=None,
                                           mpi_root=self.mpi_root,
                                           back_constraint=back_constraint, X=X,
                                           X_variance=X_variance, Z=Z,
-                                          uncertain_inputs=uncertain_inputs))
+                                          uncertain_inputs=uncertain_inputs,
+                                          psi_chunk_size=self.psi_chunk_size,
+                                          psi_threads=self.psi_threads))
         ls.append(self.layers[-1])
         self.nLayers += 1
         self.nDims.append(dim_up)