Merge pull request #21 from gsd-authors/dev

Fix imports and typing

Author: krzysztofrusek

src/gsd/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.3dev"
+__version__ = "0.2.3"
 from gsd.fit import fit_moments as fit_moments, GSDParams as GSDParams
 from gsd.gsd import (
     log_prob as log_prob,

src/gsd/gsd.py

@@ -4,7 +4,6 @@
 import jax.numpy as jnp
 import numpy as np
 from jax import Array
-from jax.random import PRNGKeyArray
 from jax.scipy.special import betaln
 from jax.typing import ArrayLike
 
@@ -113,7 +112,7 @@ def variance(psi: ArrayLike, rho: ArrayLike) -> Array:
     return rho * vmin(psi) + (1 - rho) * vmax(psi)
 
 
-def sample(psi: ArrayLike, rho: ArrayLike, shape: Shape, key: PRNGKeyArray) -> Array:
+def sample(psi: ArrayLike, rho: ArrayLike, shape: Shape, key: Array) -> Array:
     """Sample from GSD
 
     :param psi: mean

tests/experimental_test.py

@@ -16,7 +16,7 @@
 from gsd.experimental.fit import GridEstimator
 from gsd.experimental.max_entropy import MaxEntropyGSD, vmax
 from gsd.fit import fit_moments, GSDParams, log_pmax, pairs, pmax
-from gsd.gsd import make_softvmin, vmax, vmin
+from gsd.gsd import make_softvmin, vmin
 
 
 class FitTestCase(unittest.TestCase):