fix __call__ types

pyro-ppl · Dec 20, 2024 · ee126ef · ee126ef
1 parent b4dacc0
commit ee126ef
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
@@ -32,6 +32,7 @@
 import warnings
 
 import numpy as np
+from typyng import Any
 
 import jax
 from jax import lax, tree_util
@@ -572,7 +573,7 @@ class DistributionLike(Protocol):
     or tensorflow_probability.distributions.Distribution.
     """
 
-    def __call__(self, *args: functools.Any, **kwds: functools.Any) -> functools.Any:
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
         return super().__call__(*args, **kwds)
 
     @property