Make linear operators pickleable (#124)

* use qualified names as cache keys

* test pickle with populated caches

* make `KernelLinearOperator` pickle-able

Replace lambda in defaultdict with a module-level function so that
num_nonbatch_dimensions can be pickled.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: kayween

linear_operator/operators/kernel_linear_operator.py

@@ -14,6 +14,10 @@
 from linear_operator.utils.memoize import cached
 
 
+def _default_num_nonbatch_dimensions():
+    return 2
+
+
 def _x_getitem(x, batch_indices, data_index):
     """
     Helper function to compute x[*batch_indices, data_index, :] in an efficient way.
@@ -142,9 +146,9 @@ def __init__(
     ):
         # Change num_nonbatch_dimensions into a default dict
         if num_nonbatch_dimensions is None:
-            num_nonbatch_dimensions = defaultdict(lambda: 2)
+            num_nonbatch_dimensions = defaultdict(_default_num_nonbatch_dimensions)
         else:
-            num_nonbatch_dimensions = defaultdict(lambda: 2, **num_nonbatch_dimensions)
+            num_nonbatch_dimensions = defaultdict(_default_num_nonbatch_dimensions, **num_nonbatch_dimensions)
 
         # Divide params into tensors and non-tensors
         tensor_params = dict()

linear_operator/test/linear_operator_test_case.py

@@ -3,6 +3,7 @@
 
 import logging
 import math
+import pickle
 import traceback
 from abc import abstractmethod
 from itertools import combinations, product
@@ -955,6 +956,16 @@ def test_logdet(self):
             if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                 self.assertAllClose(arg.grad, arg_copy.grad, **tolerances)
 
+    def test_pickle(self):
+        linear_op = self.create_linear_op()
+
+        # Make sure that pickle works with populated caches
+        _ = linear_op.to_dense()
+
+        pickled = pickle.dumps(linear_op)
+        unpickled = pickle.loads(pickled)
+        self.assertAllClose(unpickled.to_dense(), linear_op.to_dense())
+
     def test_prod(self):
         with linear_operator.settings.fast_computations(covar_root_decomposition=False):
             linear_op = self.create_linear_op()

linear_operator/utils/memoize.py

@@ -54,7 +54,7 @@ def _cached(method=None, name=None):
 
     @functools.wraps(method)
     def g(self, *args, **kwargs):
-        cache_name = name if name is not None else method
+        cache_name = name if name is not None else method.__qualname__
         kwargs_pkl = pickle.dumps(kwargs)
         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
@@ -72,7 +72,7 @@ def _cached_ignore_args(method=None, name=None):
 
     @functools.wraps(method)
     def g(self, *args, **kwargs):
-        cache_name = name if name is not None else method
+        cache_name = name if name is not None else method.__qualname__
         if not _is_in_cache_ignore_args(self, cache_name):
             return _add_to_cache_ignore_args(self, cache_name, method(self, *args, **kwargs))
         return _get_from_cache_ignore_args(self, cache_name)