simplify array types

pyro-ppl · Dec 20, 2024 · f8ffed9 · f8ffed9
1 parent e8437bf
commit f8ffed9
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
@@ -34,7 +34,7 @@
 import numpy as np
 
 import jax
-from jax import Array, lax, tree_util
+from jax import lax, tree_util
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 from jax.typing import ArrayLike
@@ -581,7 +581,9 @@ def event_shape(self) -> tuple[int, ...]: ...
     @property
     def event_dim(self) -> int: ...
 
-    def sample(self, key: ArrayLike, sample_shape: tuple[int, ...] = ()) -> Array: ...
+    def sample(
+        self, key: ArrayLike, sample_shape: tuple[int, ...] = ()
+    ) -> ArrayLike: ...
 
     def log_prob(self, value: ArrayLike) -> ArrayLike: ...