Import Union from pysages.typing, not from plum

Author: p-zubieta

pysages/grids.py

@@ -5,9 +5,9 @@
 
 from jax import jit
 from jax import numpy as np
-from plum import Union, parametric
+from plum import parametric
 
-from pysages.typing import JaxArray
+from pysages.typing import JaxArray, Union
 from pysages.utils import dispatch, is_generic_subclass, prod