Add type hint to RandomVariable.rng_fn

aesara-devs · Oct 18, 2022 · 77df667 · 77df667
1 parent fb8224e
commit 77df667
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aesara/tensor/random/op.py b/aesara/tensor/random/op.py
@@ -161,7 +161,7 @@ def _supp_shape_from_params(self, dist_params, **kwargs):
         """
         return default_supp_shape_from_params(self.ndim_supp, dist_params, **kwargs)
 
-    def rng_fn(self, rng, *args, **kwargs):
+    def rng_fn(self, rng, *args, **kwargs) -> Union[int, float, np.ndarray]:
         """Sample a numeric random variate."""
         return getattr(rng, self.name)(*args, **kwargs)