fixing depreciation error

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = ripplegw
-version = 0.0.4
+version = 0.0.5
 description = A small jax package for differentiable and fast gravitational wave data analysis.
 long_description = file: README.md
 long_description_content_type = text/markdown

src/ripple/__init__.py

@@ -15,7 +15,7 @@
 import jax.numpy as jnp
 
 from .constants import C, G
-from .typing import Array, PRNGKeyArray
+from .typing import Array
 
 
 def Mc_eta_to_ms(m):
@@ -64,9 +64,7 @@ def get_f_isco(m):
     return 1 / (6 ** (3 / 2) * pi * m / (C**3 / G))
 
 
-def get_M_eta_sampler(
-    M_range: Tuple[float, float], eta_range: Tuple[float, float]
-) -> Callable[[PRNGKeyArray, int], Array]:
+def get_M_eta_sampler(M_range: Tuple[float, float], eta_range: Tuple[float, float]):
     """
     Uniformly values of the chirp mass and samples over the specified ranges.
     This function may be removed in the future since it is trivial.
@@ -84,9 +82,7 @@ def sampler(key, n):
     return sampler
 
 
-def get_m1_m2_sampler(
-    m1_range: Tuple[float, float], m2_range: Tuple[float, float]
-) -> Callable[[PRNGKeyArray, int], Array]:
+def get_m1_m2_sampler(m1_range: Tuple[float, float], m2_range: Tuple[float, float]):
     r"""
     Creates a function to uniformly sample two parameters, with the restriction
     that the first is larger than the second.

src/ripple/typing.py

@@ -7,5 +7,5 @@
 
 
 # TODO: what type should this be?
-PRNGKeyArray = jax.random.PRNGKeyArray  # type: ignore
+# PRNGKeyArray = jax.random.PRNGKeyArray  # type: ignore
 Array = jnp.ndarray