Use a cache on the backend itself to store the context.

uarray/_backend.py

@@ -1,24 +1,16 @@
-from typing import (
-    Callable,
-    Iterable,
-    Dict,
-    Tuple,
-    Any,
-    Set,
-    Optional,
-    Type,
-    Union,
-    List,
-)
+import typing
 import inspect
 import functools
 from . import _uarray  # type: ignore
 import copyreg  # type: ignore
 import atexit
 import pickle
+import weakref
 
-ArgumentExtractorType = Callable[..., Tuple["Dispatchable", ...]]
-ArgumentReplacerType = Callable[[Tuple, Dict, Tuple], Tuple[Tuple, Dict]]
+ArgumentExtractorType = typing.Callable[..., typing.Tuple["Dispatchable", ...]]
+ArgumentReplacerType = typing.Callable[
+    [typing.Tuple, typing.Dict, typing.Tuple], typing.Tuple[typing.Tuple, typing.Dict]
+]
 
 from ._uarray import (  # type: ignore
     BackendNotImplementedError,
@@ -103,7 +95,7 @@ def generate_multimethod(
     argument_extractor: ArgumentExtractorType,
     argument_replacer: ArgumentReplacerType,
     domain: str,
-    default: Optional[Callable] = None,
+    default: typing.Optional[typing.Callable] = None,
 ):
     """
     Generates a multimethod.
@@ -191,7 +183,16 @@ def set_backend(backend, coerce=False, only=False):
     skip_backend: A context manager that allows skipping of backends.
     set_global_backend: Set a single, global backend for a domain.
     """
-    return _SetBackendContext(backend, coerce, only)
+    try:
+        return backend.__ua_cache__["set", coerce, only]
+    except AttributeError:
+        backend.__ua_cache__ = {}
+    except KeyError:
+        pass
+
+    ctx = _SetBackendContext(backend, coerce, only)
+    backend.__ua_cache__["set", coerce, only] = ctx
+    return ctx
 
 
 def skip_backend(backend):