Merge pull request #128 from cornellius-gp/fix-deprecated-apis-and-bugs

Fix deprecated APIs, latent bugs, and code quality issues

Author: Balandat

linear_operator/operators/_linear_operator.py

@@ -80,7 +80,7 @@ def _implements_second_arg(torch_function: Callable) -> Callable:
     where the first argument of the function is a torch.Tensor and the
     second argument is a LinearOperator
 
-    Examples of this include :meth:`torch.cholesky_solve`, `torch.linalg.solve`, or `torch.matmul`.
+    Examples of this include `torch.linalg.solve` or `torch.matmul`.
     """
 
     @functools.wraps(torch_function)

linear_operator/operators/batch_repeat_linear_operator.py

@@ -139,8 +139,8 @@ def _move_repeat_batches_back(self, batch_matrix, output_shape):
         So that the tensor is now rb x m x n.
         """
         if hasattr(self, "_batch_move_memo"):
-            padded_base_batch_shape, batch_repeat = self.__batch_move_memo
-            del self.__batch_move_memo
+            _, padded_base_batch_shape, batch_repeat = self._batch_move_memo
+            del self._batch_move_memo
         else:
             padding_dims = torch.Size(tuple(1 for _ in range(max(len(output_shape) - self.base_linear_op.dim(), 0))))
             padded_base_batch_shape = padding_dims + self.base_linear_op.batch_shape
@@ -188,7 +188,7 @@ def _move_repeat_batches_to_columns(self, batch_matrix, output_shape):
         batch_matrix = batch_matrix.permute(*batch_dims, -2, -1, *repeat_dims).contiguous()
         batch_matrix = batch_matrix.view(*self.base_linear_op.batch_shape, output_shape[-2], -1)
 
-        self.__batch_move_memo = output_shape, padded_base_batch_shape, batch_repeat
+        self._batch_move_memo = output_shape, padded_base_batch_shape, batch_repeat
         return batch_matrix
 
     def _permute_batch(self, *dims: int) -> LinearOperator:

linear_operator/operators/cat_linear_operator.py

@@ -11,7 +11,6 @@
 from linear_operator.operators.dense_linear_operator import DenseLinearOperator, to_linear_operator
 
 from linear_operator.utils.broadcasting import _matmul_broadcast_shape
-from linear_operator.utils.deprecation import bool_compat
 from linear_operator.utils.generic import _to_helper
 from linear_operator.utils.getitem import _noop_index
 
@@ -188,7 +187,7 @@ def _get_indices(self, row_index: IndexType, col_index: IndexType, *batch_indice
 
         # Find out for which indices we switch to different tensors
         target_tensors = self.idx_to_tensor_idx[cat_dim_indices]
-        does_switch_tensor = torch.ones(target_tensors.numel() + 1, dtype=bool_compat, device=self.device)
+        does_switch_tensor = torch.ones(target_tensors.numel() + 1, dtype=torch.bool, device=self.device)
         torch.ne(target_tensors[:-1], target_tensors[1:], out=does_switch_tensor[1:-1])
 
         # Get the LinearOperators that will comprise the new LinearOperator
@@ -258,7 +257,7 @@ def _getitem(self, row_index: IndexType, col_index: IndexType, *batch_indices: I
 
             # Find out for which indices we switch to different tensors
             target_tensors = self.idx_to_tensor_idx[cat_dim_indices]
-            does_switch_tensor = torch.ones(target_tensors.numel() + 1, dtype=bool_compat, device=self.device)
+            does_switch_tensor = torch.ones(target_tensors.numel() + 1, dtype=torch.bool, device=self.device)
             torch.ne(target_tensors[:-1], target_tensors[1:], out=does_switch_tensor[1:-1])
 
             # Get the LinearOperators that will comprise the new LinearOperator
@@ -294,7 +293,7 @@ def _getitem(self, row_index: IndexType, col_index: IndexType, *batch_indices: I
 
         else:
             raise RuntimeError(
-                "Unexpected index type {cat_dim_indices.__class__.__name__}. This is a bug in LinearOperator."
+                f"Unexpected index type {cat_dim_indices.__class__.__name__}. This is a bug in LinearOperator."
             )
 
         # Process the list

linear_operator/operators/kronecker_product_linear_operator.py

@@ -35,7 +35,7 @@ def _matmul(linear_ops, kp_shape, rhs):
     output_shape = _matmul_broadcast_shape(kp_shape, rhs.shape)
     output_batch_shape = output_shape[:-2]
 
-    res = rhs.contiguous().expand(*output_batch_shape, *rhs.shape[-2:])
+    res = rhs.expand(*output_batch_shape, *rhs.shape[-2:])
     num_cols = rhs.size(-1)
     for linear_op in linear_ops:
         res = res.view(*output_batch_shape, linear_op.size(-1), -1)
@@ -50,7 +50,7 @@ def _t_matmul(linear_ops, kp_shape, rhs):
     output_shape = _matmul_broadcast_shape(kp_t_shape, rhs.shape)
     output_batch_shape = torch.Size(output_shape[:-2])
 
-    res = rhs.contiguous().expand(*output_batch_shape, *rhs.shape[-2:])
+    res = rhs.expand(*output_batch_shape, *rhs.shape[-2:])
     num_cols = rhs.size(-1)
     for linear_op in linear_ops:
         res = res.view(*output_batch_shape, linear_op.size(-2), -1)

linear_operator/operators/low_rank_root_added_diag_linear_operator.py

@@ -77,7 +77,9 @@ def _solve(
         chol_cap_mat = self.chol_cap_mat
 
         res = V.matmul(A_inv.matmul(rhs))
-        res = torch.cholesky_solve(res, chol_cap_mat)
+        res = torch.linalg.solve_triangular(
+            chol_cap_mat.mT, torch.linalg.solve_triangular(chol_cap_mat, res, upper=False), upper=True
+        )
         res = A_inv.matmul(U.matmul(res))
 
         solve = A_inv.matmul(rhs) - res

linear_operator/operators/sum_linear_operator.py

@@ -28,7 +28,7 @@ def __init__(self, *linear_ops, **kwargs):
     def _diagonal(
         self: LinearOperator,  # shape: (..., M, N)
     ) -> torch.Tensor:  # shape: (..., N)
-        return sum(linear_op._diagonal().contiguous() for linear_op in self.linear_ops)
+        return sum(linear_op._diagonal() for linear_op in self.linear_ops)
 
     def _expand_batch(
         self: LinearOperator, batch_shape: torch.Size | list[int]  # shape: (..., M, N)

linear_operator/operators/triangular_linear_operator.py

@@ -194,7 +194,7 @@ def inv_quad_logdet(
             inv_quad_term = (inv_quad_rhs * self.solve(inv_quad_rhs)).sum(dim=-2)
         if logdet:
             diag = self._diagonal()
-            logdet_term = self._diagonal().abs().log().sum(-1)
+            logdet_term = diag.abs().log().sum(-1)
             if torch.sign(diag).prod(-1) < 0:
                 logdet_term = torch.full_like(logdet_term, float("nan"))
         else:

linear_operator/test/linear_operator_test_case.py

@@ -666,7 +666,6 @@ def test_add_jitter(self):
         self.assertAllClose(res, actual)
 
     def test_add_low_rank(self):
-        linear_op = self.create_linear_op()
         linear_op = self.create_linear_op()
         evaluated = self.evaluate_linear_op(linear_op)
         new_rows = torch.randn(*linear_op.shape[:-1], 3)

linear_operator/utils/deprecation.py

@@ -3,15 +3,6 @@
 
 import functools
 import warnings
-from unittest.mock import MagicMock
-
-import torch
-
-# TODO: Use bool instead of uint8 dtype once pytorch #21113 is in stable release
-if isinstance(torch, MagicMock):
-    bool_compat = torch.uint8
-else:
-    bool_compat = (torch.ones(1) > 0).dtype
 
 
 class DeprecationError(Exception):

linear_operator/utils/linear_cg.py

@@ -6,7 +6,6 @@
 import torch
 
 from linear_operator import settings
-from linear_operator.utils.deprecation import bool_compat
 from linear_operator.utils.warnings import NumericalWarning
 
 
@@ -219,7 +218,7 @@ def linear_cg(
         mul_storage = torch.empty_like(residual)
         alpha = torch.empty(*batch_shape, 1, rhs.size(-1), dtype=residual.dtype, device=residual.device)
         beta = torch.empty_like(alpha)
-        is_zero = torch.empty(*batch_shape, 1, rhs.size(-1), dtype=bool_compat, device=residual.device)
+        is_zero = torch.empty(*batch_shape, 1, rhs.size(-1), dtype=torch.bool, device=residual.device)
 
     # Define tridiagonal matrices, if applicable
     if n_tridiag:
@@ -231,7 +230,7 @@ def linear_cg(
             dtype=alpha.dtype,
             device=alpha.device,
         )
-        alpha_tridiag_is_zero = torch.empty(*batch_shape, n_tridiag, dtype=bool_compat, device=t_mat.device)
+        alpha_tridiag_is_zero = torch.empty(*batch_shape, n_tridiag, dtype=torch.bool, device=t_mat.device)
         alpha_reciprocal = torch.empty(*batch_shape, n_tridiag, dtype=t_mat.dtype, device=t_mat.device)
         prev_alpha_reciprocal = torch.empty_like(alpha_reciprocal)
         prev_beta = torch.empty_like(alpha_reciprocal)