LazyTensor -> LinearOperator (#2027)

* Bump version, require linear_operator

* [WIP] Replace LazyTensor with LinearOperator, fix gpytorch/lazy test suite failures

* [WIP] fix gpytorch/variational test suite failures

* [WIP] fix gpytorch/functions test suite failures

* [WIP] specify linear_operator commit in setup.py

* [WIP] Remove lazy docs

* [WIP] Update CI requirements

* [WIP] Pipe eval_cg_tolerance to linear_operator.utils.linear_cg

* [WIP] Use linear_operator errors (fix variational fantasy test failures)

* [WIP] fix remaining test suite bugs

* [WIP] Use linear_operator functions

* [WIP] Add linear_operator dependency, bump Python to 3.8

* [WIP] remove unused utilities

* [WIP] Move Cholesky to LinearOperator

* Fix rebase issues

* Fix new failing test, consolodate requirements, fix docs typo.

* Fix import in IdentityLT test

* Bump linear operator version

* Fix cuda issue in polya gamma notebook

* Shorten LinearOperator and Tensor imports in gpytorch/__init__.py

* Remove deprecation warnings after linear_operator migration (#2102)

* Remove tests of deprecated items that were moved to linear_operator

* Run test suite

* Replace all LazyTensors with LinearOperators

* Update .diag, .evaluate, .add_diag, .inv_matmul calls, solve ConstantMean deprecation warnings

* Rename *_lazy_tensor arguments, use linear_operator interpolation utils

* Resolve solve_triangular warnings

* Resolve eigh warnings

* Resolve meshgrid warnings

* Update gpytorch/kernels/multi_device_kernel.py

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>

* Update gpytorch/likelihoods/multitask_gaussian_likelihood.py

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>

* Consistent  import

* Add tests for deprecated LazyTensor import

* Fix doc issues

* Fix doc issues

* Remove extra testing requirements

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>

Author: gpleiss

.conda/meta.yaml

@@ -13,13 +13,12 @@ build:
 
 requirements:
   host:
-    - python>=3.7
+    - python>=3.8
 
   run:
-    - pytorch>=1.10
-    - numpy
+    - pytorch>=1.11
     - scikit-learn
-    - scipy
+    - linear_operator>=0.1.1
 
 test:
   imports:

.github/workflows/deploy.yml

@@ -19,7 +19,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: "3.7"
+        python-version: "3.8"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -40,7 +40,7 @@ jobs:
     - uses: conda-incubator/setup-miniconda@v2
       with:
         auto-update-conda: false
-        python-version: "3.7"
+        python-version: "3.8"
     - name: Install dependencies
       run: |
         conda install -y anaconda-client conda-build

.github/workflows/run_test_suite.yml

@@ -18,7 +18,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: "3.7"
+        python-version: "3.8"
     - name: Install dependencies
       run: |
         pip install flake8==4.0.1 flake8-print==4.0.0 pre-commit
@@ -43,13 +43,13 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: "3.7"
+        python-version: "3.8"
     - name: Install dependencies
       run: |
         if [[ ${{ matrix.pytorch-version }} = "master" ]]; then
           pip install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html;
         else
-          pip install torch==1.10+cpu -f https://download.pytorch.org/whl/torch_stable.html;
+          pip install torch==1.11+cpu -f https://download.pytorch.org/whl/torch_stable.html;
         fi
         if [[ ${{ matrix.extras }} == "with-extras" ]]; then
           pip install "pyro-ppl>=1.8";
@@ -67,9 +67,10 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: "3.7"
+        python-version: "3.8"
     - name: Install dependencies
       run: |
+        pip install torch==1.11+cpu -f https://download.pytorch.org/whl/torch_stable.html
         pip install pytest nbval jupyter tqdm matplotlib torchvision scipy
         pip install -r requirements.txt
         python setup.py build develop

.readthedocs.yml

@@ -16,7 +16,7 @@ sphinx:
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-  version: 3.7
+  version: 3.8
   install:
     - requirements: requirements.txt
     - requirements: docs/requirements.txt

README.md

@@ -6,7 +6,7 @@
 
 GPyTorch is a Gaussian process library implemented using PyTorch. GPyTorch is designed for creating scalable, flexible, and modular Gaussian process models with ease.
 
-Internally, GPyTorch differs from many existing approaches to GP inference by performing all inference operations using modern numerical linear algebra techniques like preconditioned conjugate gradients. Implementing a scalable GP method is as simple as providing a matrix multiplication routine with the kernel matrix and its derivative via our `LazyTensor` interface, or by composing many of our already existing `LazyTensors`. This allows not only for easy implementation of popular scalable GP techniques, but often also for significantly improved utilization of GPU computing compared to solvers based on the Cholesky decomposition.
+Internally, GPyTorch differs from many existing approaches to GP inference by performing all inference operations using modern numerical linear algebra techniques like preconditioned conjugate gradients. Implementing a scalable GP method is as simple as providing a matrix multiplication routine with the kernel matrix and its derivative via our `LinearOperator` interface, or by composing many of our already existing `LinearOperators`. This allows not only for easy implementation of popular scalable GP techniques, but often also for significantly improved utilization of GPU computing compared to solvers based on the Cholesky decomposition.
 
 GPyTorch provides (1) significant GPU acceleration (through MVM based inference); (2) state-of-the-art implementations of the latest algorithmic advances for scalability and flexibility ([SKI/KISS-GP](http://proceedings.mlr.press/v37/wilson15.pdf), [stochastic Lanczos expansions](https://arxiv.org/abs/1711.03481), [LOVE](https://arxiv.org/pdf/1803.06058.pdf), [SKIP](https://arxiv.org/pdf/1802.08903.pdf), [stochastic variational](https://arxiv.org/pdf/1611.00336.pdf) [deep kernel learning](http://proceedings.mlr.press/v51/wilson16.pdf), ...); (3) easy integration with deep learning frameworks.
 
@@ -17,8 +17,8 @@ See our numerous [**examples and tutorials**](https://gpytorch.readthedocs.io/en
 ## Installation
 
 **Requirements**:
-- Python >= 3.7
-- PyTorch >= 1.10
+- Python >= 3.8
+- PyTorch >= 1.11
 
 Install GPyTorch using pip or conda:
 

docs/source/distributions.rst

@@ -9,7 +9,7 @@ For the most part, GpyTorch relies on torch's distribution library.
 However, we offer two custom distributions.
 
 We implement a custom :obj:`~gpytorch.distributions.MultivariateNormal` that accepts
-:obj:`~gpytorch.lazy.LazyTensor` objects for covariance matrices. This allows us to use custom
+:obj:`~linear_operator.operators.LinearOperator` objects for covariance matrices. This allows us to use custom
 linear algebra operations, which makes this more efficient than PyTorch's MVN implementation.
 
 In addition, we implement a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal` which

docs/source/functions.rst

@@ -12,26 +12,24 @@ Functions
 
 .. automodule:: gpytorch
 
-.. autofunction:: add_diag
+.. autofunction:: add_diagonal
 
 .. autofunction:: add_jitter
 
 .. autofunction:: dsmm
 
-.. autofunction:: inv_matmul
+.. autofunction:: diagonalization
 
 .. autofunction:: inv_quad
 
 .. autofunction:: inv_quad_logdet
 
-.. autofunction:: matmul
-
-.. autofunction:: logdet
-
-.. autofunction:: log_normal_cdf
-
 .. autofunction:: pivoted_cholesky
 
 .. autofunction:: root_decomposition
 
 .. autofunction:: root_inv_decomposition
+
+.. autofunction:: solve
+
+.. autofunction:: sqrt_inv_matmul

docs/source/index.rst

@@ -50,7 +50,6 @@ GPyTorch's documentation
    :caption: Advanced Package Reference
 
    module
-   lazy
    functions
    utils
 

docs/source/lazy.rst

@@ -5,42 +5,31 @@ gpytorch.utils
 ===================================
 
 .. currentmodule:: gpytorch.utils
-
-
-Utilities
-----------------
-
 .. automodule:: gpytorch.utils
    :members:
 
-Lanczos Utilities
-~~~~~~~~~~~~~~~~~
+Grid Utilities
+----------------------
 
-.. automodule:: gpytorch.utils.lanczos
+.. automodule:: gpytorch.utils.grid
    :members:
 
-Permutation Utilities
-~~~~~~~~~~~~~~~~~~~~~~~~~~
+Interpolation Utilities
+---------------------------------
 
-.. automodule:: gpytorch.utils.permutation
+.. automodule:: gpytorch.utils.interpolation
    :members:
 
 Quadrature Utilities
-~~~~~~~~~~~~~~~~~~~~
+----------------------------
 
 .. automodule:: gpytorch.utils.quadrature
    :members:
 
-Sparse Utilities
-~~~~~~~~~~~~~~~~~
+Transform Utilities
+--------------------------
 
-.. automodule:: gpytorch.utils.sparse
-   :members:
-
-Grid Utilities
-~~~~~~~~~~~~~~~~~
-
-.. automodule:: gpytorch.utils.grid
+.. automodule:: gpytorch.utils.transforms
    :members:
 
 Nearest Neighbors Utilities

docs/source/utils.rst

@@ -1,8 +1,8 @@
 name: gpytorch
 channels:
   - pytorch
+  - gpytorch
 dependencies:
   - pytorch
-  - numpy
   - scikit-learn
-  - scipy
+  - linear_operator>0.1.1

environment.yml

@@ -96,7 +96,7 @@
     "\n",
     "To implement a custom kernel, we derive one from GPyTorch's [kernel class](https://docs.gpytorch.ai/en/latest/kernels.html) and implement the `forward()` method. The base class provides many useful routines. For example, `__call__()` is implemented, so that the kernel may be called directly, without resorting to the `forward()` routine. Among other things, the `Kernel` class provides a method `covar_dist()`, which may be used to calculate the Euclidian distance between point pairs conveniently.\n",
     "\n",
-    "The `forward()` method represents the kernel function and should return a `torch.tensor` or a `gpytorch.lazy.LazyTensor`, when called on two `torch.tensor`s:"
+    "The `forward()` method represents the kernel function and should return a `torch.tensor` or a `linear_operator.operators.LinearOperator`, when called on two `torch.tensor`s:"
    ]
   },
   {

examples/00_Basic_Usage/Implementing_a_custom_Kernel.ipynb

@@ -128,7 +128,7 @@
     }
    ],
    "source": [
-    "plt.plot(train_x.squeeze(-1), train_y, 'o')"
+    "plt.plot(train_x.squeeze(-1).cpu(), train_y.cpu(), 'o')"
    ]
   },
   {
@@ -1755,7 +1755,7 @@
     "# push training data points through model\n",
     "train_mean_f = model(train_x).loc.data.cpu()\n",
     "# plot training data with y being -1/1 valued\n",
-    "plt.plot(train_x.squeeze(-1), train_y.mul(2.).sub(1.), 'o')\n",
+    "plt.plot(train_x.squeeze(-1).cpu(), train_y.mul(2.).sub(1.).cpu(), 'o')\n",
     "# plot mean gaussian process posterior mean evaluated at training data\n",
     "plt.plot(train_x.squeeze(-1).cpu(), train_mean_f.cpu(), 'x')"
    ]

examples/04_Variational_and_Approximate_GPs/PolyaGamma_Binary_Classification.ipynb

@@ -71,7 +71,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "xv, yv = torch.meshgrid([torch.linspace(0, 1, 10), torch.linspace(0, 1, 10)])\n",
+    "xv, yv = torch.meshgrid(torch.linspace(0, 1, 10), torch.linspace(0, 1, 10), indexing=\"ij\")\n",
     "train_x = torch.cat((\n",
     "    xv.contiguous().view(xv.numel(), 1), \n",
     "    yv.contiguous().view(yv.numel(), 1)),\n",
@@ -301,7 +301,7 @@
     "\n",
     "# Test points\n",
     "n1, n2 = 50, 50\n",
-    "xv, yv = torch.meshgrid([torch.linspace(0, 1, n1), torch.linspace(0, 1, n2)])\n",
+    "xv, yv = torch.meshgrid(torch.linspace(0, 1, n1), torch.linspace(0, 1, n2), indexing=\"ij\")\n",
     "f, dfx, dfy = franke(xv, yv)\n",
     "\n",
     "# Make predictions\n",