Use explicit Union rather than | in type hints for Python 3.9 compat

Author: matt-graham

s2fft/utils/jax_primitive.py

@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Callable
+from typing import Callable, Dict, Optional, Union
 from jax import core
 from jax.interpreters import ad, batching, xla, mlir
 
@@ -8,10 +8,10 @@ def register_primitive(
     name: str,
     multiple_results: bool,
     abstract_evaluation: Callable,
-    lowering_per_platform: dict[None | str, Callable],
-    batcher: None | Callable = None,
-    jacobian_vector_product: None | Callable = None,
-    transpose: None | Callable = None,
+    lowering_per_platform: Dict[Union[None, str], Callable],
+    batcher: Optional[Callable] = None,
+    jacobian_vector_product: Optional[Callable] = None,
+    transpose: Optional[Callable] = None,
 ):
     """Register a new custom JAX primitive.