Plum 2 compatibility

setup.py

@@ -4,7 +4,7 @@
     "numpy>=1.16",
     "fdm",
     "algebra>=1",
-    "plum-dispatch>=1.5.3",
+    "plum-dispatch>=2",
     "backends>=1.4.11",
     "backends-matrix>=1.2.11",
     "mlkernels>=0.3.6",

stheno/model/gp.py

@@ -1,17 +1,14 @@
 from types import FunctionType
+from typing import Union
 
 from fdm import central_fdm
 from lab import B
-from mlkernels import (
-    OneKernel,
-    ZeroMean,
-    OneMean,
-)
-from plum import Union
+from mlkernels import OneKernel, OneMean, ZeroMean
+from plum import isinstance
 
-from .fdd import FDD
 from .. import PromisedGP, PromisedMeasure, _dispatch
 from ..random import RandomProcess
+from .fdd import FDD
 
 __all__ = ["assert_same_measure", "intersection_measure_group", "cross", "GP"]
 

stheno/model/measure.py

@@ -1,29 +1,23 @@
 from types import FunctionType
+from typing import Union
 
 from lab import B
 from matrix import Constant
-from mlkernels import (
-    num_elements,
-    ZeroKernel,
-    TensorProductKernel,
-)
-from plum import Union
+from mlkernels import TensorProductKernel, ZeroKernel, num_elements
 
+from .. import PromisedMeasure, _dispatch
+from ..lazy import LazyMatrix, LazyVector
+from ..mo import AmbiguousDimensionalityKernel as ADK
+from ..mo import MultiOutputKernel as MOK
+from ..mo import MultiOutputMean as MOM
 from .fdd import FDD
 from .gp import GP, assert_same_measure
 from .observations import (
     AbstractObservations,
-    Observations,
     AbstractPseudoObservations,
+    Observations,
     combine,
 )
-from .. import _dispatch, PromisedMeasure
-from ..lazy import LazyVector, LazyMatrix
-from ..mo import (
-    MultiOutputKernel as MOK,
-    MultiOutputMean as MOM,
-    AmbiguousDimensionalityKernel as ADK,
-)
 
 __all__ = ["Measure"]
 

stheno/model/observations.py

@@ -1,7 +1,8 @@
+from typing import Union
+
 from lab import B
 from matrix import AbstractMatrix, Diagonal
 from mlkernels import PosteriorKernel, PosteriorMean, SubspaceKernel, num_elements
-from plum import Union
 
 from .. import _dispatch
 from .fdd import FDD

stheno/random.py

@@ -1,9 +1,10 @@
 from types import FunctionType
+from typing import Union
 
 from algebra.util import identical
 from lab import B
-from matrix import AbstractMatrix, Zero, Diagonal
-from plum import convert, Union
+from matrix import AbstractMatrix, Diagonal, Zero
+from plum import convert
 from wbml.util import indented_kv
 
 from . import _dispatch

tests/model/test_model.py

@@ -3,28 +3,21 @@
 import tensorflow as tf
 from lab import B
 from matrix import Diagonal
-from mlkernels import (
-    Kernel,
-    pairwise,
-    elwise,
-    Linear,
-    EQ,
-    Delta,
-    Exp,
-)
+from mlkernels import EQ, Delta, Exp, Kernel, Linear, elwise, pairwise
+from plum import isinstance
 from stheno.model import (
-    Measure,
+    FDD,
     GP,
+    Measure,
     Obs,
     PseudoObs,
-    PseudoObsFITC,
     PseudoObsDTC,
+    PseudoObsFITC,
     cross,
-    FDD,
 )
 
-from .util import assert_equal_normals, assert_equal_measures
 from ..util import approx
+from .util import assert_equal_measures, assert_equal_normals
 
 
 def test_measure_groups():

tests/test_random.py

@@ -3,7 +3,7 @@
 import pytest
 from algebra.util import identical
 from matrix import Dense, Zero
-from plum import NotFoundLookupError
+from plum import NotFoundLookupError, isinstance
 from scipy.stats import multivariate_normal
 from stheno.random import Normal, RandomVector
 

tests/util.py

@@ -1,10 +1,10 @@
 from time import time
+from typing import Union
 
 import numpy as np
 from lab import B
 from matrix import AbstractMatrix
-from plum import dispatch, Union
-
+from plum import dispatch
 from stheno import Normal
 
 __all__ = ["benchmark", "approx"]