fix mypy errors

Author: kratsg

src/pyhf/tensor/numpy_backend.py

@@ -18,6 +18,7 @@
 from scipy.stats import norm, poisson
 
 from pyhf.typing import Literal, Shape
+from typing import cast
 
 T = TypeVar("T", bound=NBitBase)
 
@@ -26,27 +27,32 @@
 log = logging.getLogger(__name__)
 
 
-class _BasicPoisson:
+class _BasicPoisson(Generic[T]):
     def __init__(self, rate: Tensor[T]):
         self.rate = rate
 
     def sample(self, sample_shape: Shape) -> ArrayLike:
-        return poisson(self.rate).rvs(size=sample_shape + self.rate.shape)  # type: ignore[no-any-return]
+        return cast(
+            ArrayLike, poisson(self.rate).rvs(size=sample_shape + self.rate.shape)
+        )
 
-    def log_prob(self, value: NDArray[np.number[T]]) -> ArrayLike:
+    def log_prob(self, value: Tensor[T]) -> ArrayLike:
         tensorlib: numpy_backend[T] = numpy_backend()
         return tensorlib.poisson_logpdf(value, self.rate)
 
 
-class _BasicNormal:
+class _BasicNormal(Generic[T]):
     def __init__(self, loc: Tensor[T], scale: Tensor[T]):
         self.loc = loc
         self.scale = scale
 
     def sample(self, sample_shape: Shape) -> ArrayLike:
-        return norm(self.loc, self.scale).rvs(size=sample_shape + self.loc.shape)  # type: ignore[no-any-return]
+        return cast(
+            ArrayLike,
+            norm(self.loc, self.scale).rvs(size=sample_shape + self.loc.shape),
+        )
 
-    def log_prob(self, value: NDArray[np.number[T]]) -> ArrayLike:
+    def log_prob(self, value: Tensor[T]) -> ArrayLike:
         tensorlib: numpy_backend[T] = numpy_backend()
         return tensorlib.normal_logpdf(value, self.loc, self.scale)
 
@@ -125,7 +131,7 @@ def erf(self, tensor_in: Tensor[T]) -> ArrayLike:
         Returns:
             NumPy ndarray: The values of the error function at the given points.
         """
-        return special.erf(tensor_in)  # type: ignore[no-any-return]
+        return cast(ArrayLike, special.erf(tensor_in))
 
     def erfinv(self, tensor_in: Tensor[T]) -> ArrayLike:
         """
@@ -145,7 +151,7 @@ def erfinv(self, tensor_in: Tensor[T]) -> ArrayLike:
         Returns:
             NumPy ndarray: The values of the inverse of the error function at the given points.
         """
-        return special.erfinv(tensor_in)  # type: ignore[no-any-return]
+        return cast(ArrayLike, special.erfinv(tensor_in))
 
     def tile(self, tensor_in: Tensor[T], repeats: int | Sequence[int]) -> ArrayLike:
         """
@@ -207,7 +213,7 @@ def tolist(self, tensor_in: Tensor[T] | list[T]) -> list[T]:
             raise
 
     def outer(self, tensor_in_1: Tensor[T], tensor_in_2: Tensor[T]) -> ArrayLike:
-        return np.outer(tensor_in_1, tensor_in_2)  # type: ignore[arg-type]
+        return cast(ArrayLike, np.outer(tensor_in_1, tensor_in_2))
 
     def gather(self, tensor: Tensor[T], indices: NDArray[np.integer[T]]) -> ArrayLike:
         return tensor[indices]
@@ -255,7 +261,7 @@ def sum(self, tensor_in: Tensor[T], axis: int | None = None) -> ArrayLike:
         return np.sum(tensor_in, axis=axis)
 
     def product(self, tensor_in: Tensor[T], axis: Shape | None = None) -> ArrayLike:
-        return np.prod(tensor_in, axis=axis)  # type: ignore[arg-type]
+        return cast(ArrayLike, np.prod(tensor_in, axis=axis))
 
     def abs(self, tensor: Tensor[T]) -> ArrayLike:
         return np.abs(tensor)
@@ -345,7 +351,7 @@ def percentile(
         .. versionadded:: 0.7.0
         """
         # see https://github.com/numpy/numpy/issues/22125
-        return np.percentile(tensor_in, q, axis=axis, interpolation=interpolation)  # type: ignore[call-overload,no-any-return]
+        return cast(ArrayLike, np.percentile(tensor_in, q, axis=axis, interpolation=interpolation))  # type: ignore[call-overload]
 
     def stack(self, sequence: Sequence[Tensor[T]], axis: int = 0) -> ArrayLike:
         return np.stack(sequence, axis=axis)
@@ -392,7 +398,7 @@ def simple_broadcast(self, *args: Sequence[Tensor[T]]) -> Sequence[Tensor[T]]:
         return np.broadcast_arrays(*args)
 
     def shape(self, tensor: Tensor[T]) -> Shape:
-        return tensor.shape
+        return cast(Shape, tensor.shape)
 
     def reshape(self, tensor: Tensor[T], newshape: Shape) -> ArrayLike:
         return np.reshape(tensor, newshape)
@@ -434,10 +440,10 @@ def einsum(self, subscripts: str, *operands: Sequence[Tensor[T]]) -> ArrayLike:
         Returns:
             tensor: the calculation based on the Einstein summation convention
         """
-        return np.einsum(subscripts, *operands)  # type: ignore[arg-type,no-any-return]
+        return cast(ArrayLike, np.einsum(subscripts, *operands))
 
     def poisson_logpdf(self, n: Tensor[T], lam: Tensor[T]) -> ArrayLike:
-        return xlogy(n, lam) - lam - gammaln(n + 1.0)  # type: ignore[no-any-return]
+        return cast(ArrayLike, xlogy(n, lam) - lam - gammaln(n + 1.0))
 
     def poisson(self, n: Tensor[T], lam: Tensor[T]) -> ArrayLike:
         r"""
@@ -481,7 +487,7 @@ def poisson(self, n: Tensor[T], lam: Tensor[T]) -> ArrayLike:
         """
         _n = np.asarray(n)
         _lam = np.asarray(lam)
-        return np.exp(xlogy(_n, _lam) - _lam - gammaln(_n + 1.0))  # type: ignore[no-any-return,operator]
+        return cast(ArrayLike, np.exp(xlogy(_n, _lam) - _lam - gammaln(_n + 1)))
 
     def normal_logpdf(self, x: Tensor[T], mu: Tensor[T], sigma: Tensor[T]) -> ArrayLike:
         # this is much faster than
@@ -491,7 +497,7 @@ def normal_logpdf(self, x: Tensor[T], mu: Tensor[T], sigma: Tensor[T]) -> ArrayL
         root2pi = np.sqrt(2 * np.pi)
         prefactor = -np.log(sigma * root2pi)
         summand = -np.square(np.divide((x - mu), (root2 * sigma)))
-        return prefactor + summand  # type: ignore[no-any-return]
+        return cast(ArrayLike, prefactor + summand)
 
     # def normal_logpdf(self, x, mu, sigma):
     #     return norm.logpdf(x, loc=mu, scale=sigma)
@@ -522,7 +528,7 @@ def normal(self, x: Tensor[T], mu: Tensor[T], sigma: Tensor[T]) -> ArrayLike:
         Returns:
             NumPy float: Value of Normal(x|mu, sigma)
         """
-        return norm.pdf(x, loc=mu, scale=sigma)  # type: ignore[no-any-return]
+        return cast(ArrayLike, norm.pdf(x, loc=mu, scale=sigma))
 
     def normal_cdf(
         self, x: Tensor[T], mu: float | Tensor[T] = 0, sigma: float | Tensor[T] = 1
@@ -548,9 +554,9 @@ def normal_cdf(
         Returns:
             NumPy float: The CDF
         """
-        return norm.cdf(x, loc=mu, scale=sigma)  # type: ignore[no-any-return]
+        return cast(ArrayLike, norm.cdf(x, loc=mu, scale=sigma))
 
-    def poisson_dist(self, rate: Tensor[T]) -> _BasicPoisson:
+    def poisson_dist(self, rate: Tensor[T]) -> _BasicPoisson[T]:
         r"""
         The Poisson distribution with rate parameter :code:`rate`.
 
@@ -571,7 +577,7 @@ def poisson_dist(self, rate: Tensor[T]) -> _BasicPoisson:
         """
         return _BasicPoisson(rate)
 
-    def normal_dist(self, mu: Tensor[T], sigma: Tensor[T]) -> _BasicNormal:
+    def normal_dist(self, mu: Tensor[T], sigma: Tensor[T]) -> _BasicNormal[T]:
         r"""
         The Normal distribution with mean :code:`mu` and standard deviation :code:`sigma`.