Revert dual-backend Constraint.__call__ returns to ArrayLike

The np/jnp constraint methods (_CorrCholesky, _LowerCholesky, _L1Ball,
_PositiveDefinite, _PositiveSemiDefinite, _IndependentConstraint) can
return a numpy array on the numpy backend, so ty infers an
`np.ndarray | Array` union that the narrowed `-> Array` annotation
rejects. A blanket `# type: ignore` was unstable across numpy stub
versions: needed under numpy 2.5.0 (py3.14) but flagged unused under
numpy 2.4.6 (py3.11), failing lint on one Python version or the other.

Revert these overrides plus the base Constraint.__call__/check and the
two _validate_sample passthroughs to ArrayLike, which typechecks under
both numpy versions. The ~20 pure-jnp constraint overrides keep their
narrowed -> Array (a valid covariant override).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: tillahoffmann

numpyro/distributions/constraints.py

@@ -97,13 +97,13 @@ def __init_subclass__(cls, **kwargs):
     def tree_flatten(self):
         raise NotImplementedError
 
-    def __call__(self, x: NumLikeT) -> Array:
+    def __call__(self, x: NumLikeT) -> ArrayLike:
         raise NotImplementedError
 
     def __repr__(self) -> str:
         return self.__class__.__name__[1:] + "()"
 
-    def check(self, value: NumLikeT) -> Array:
+    def check(self, value: NumLikeT) -> ArrayLike:
         """
         Returns a byte tensor of `sample_shape + batch_shape` indicating
         whether each event in value satisfies this constraint.
@@ -171,7 +171,7 @@ def feasible_like(self, prototype: NumLike) -> NumLike:
 class _CorrCholesky(_SingletonConstraint[NonScalarArray]):
     event_dim = 2
 
-    def __call__(self, x: NonScalarArray) -> Array:
+    def __call__(self, x: NonScalarArray) -> ArrayLike:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         tril = xp.tril(x)
         lower_triangular = xp.all(xp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)
@@ -387,7 +387,7 @@ def is_discrete(self) -> bool:
     def event_dim(self) -> int:
         return self.base_constraint.event_dim + self.reinterpreted_batch_ndims
 
-    def __call__(self, x: NumLikeT) -> Array:
+    def __call__(self, x: NumLikeT) -> ArrayLike:
         result = self.base_constraint(x)
         if self.reinterpreted_batch_ndims == 0:
             return result
@@ -641,7 +641,7 @@ def __repr__(self) -> str:
 class _LowerCholesky(_SingletonConstraint[NonScalarArray]):
     event_dim = 2
 
-    def __call__(self, x: NonScalarArray) -> Array:
+    def __call__(self, x: NonScalarArray) -> ArrayLike:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         tril = xp.tril(x)
         lower_triangular = xp.all(xp.reshape(tril == x, x.shape[:-2] + (-1,)), axis=-1)
@@ -687,7 +687,7 @@ class _L1Ball(_SingletonConstraint[NumLike]):
     event_dim = 1
     reltol = 10.0  # Relative to finfo.eps.
 
-    def __call__(self, x: NumLike) -> Array:
+    def __call__(self, x: NumLike) -> ArrayLike:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         dtype = x.dtype if isinstance(x, xp.ndarray) else type(x)
         eps = jnp.finfo(dtype).eps
@@ -710,7 +710,7 @@ def feasible_like(self, prototype: NonScalarArray) -> NonScalarArray:
 class _PositiveDefinite(_SingletonConstraint[NonScalarArray]):
     event_dim = 2
 
-    def __call__(self, x: NonScalarArray) -> Array:
+    def __call__(self, x: NonScalarArray) -> ArrayLike:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         # check for symmetric
         symmetric = xp.all(xp.isclose(x, xp.swapaxes(x, -2, -1)), axis=(-2, -1))
@@ -738,7 +738,7 @@ def feasible_like(self, prototype: NonScalarArray) -> NonScalarArray:
 class _PositiveSemiDefinite(_SingletonConstraint[NonScalarArray]):
     event_dim = 2
 
-    def __call__(self, x: NonScalarArray) -> Array:
+    def __call__(self, x: NonScalarArray) -> ArrayLike:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         # check for symmetric
         symmetric = xp.all(xp.isclose(x, xp.swapaxes(x, -2, -1)), axis=(-2, -1))
@@ -853,11 +853,10 @@ def event_dim(self) -> int:
     def __call__(self, x: NonScalarArray) -> Array:
         xp = np if isinstance(x, (np.ndarray, np.generic)) else jnp
         tol = xp.finfo(x.dtype).eps * x.shape[-1] * 10
-        zerosum_true = True
+        zerosum_true = jnp.asarray(True)
         for dim in range(-self.event_dim, 0):
             zerosum_true = zerosum_true & xp.allclose(x.sum(dim), 0, atol=tol)
-        # FIXME: shape must match batch shape of `x`, not be a literal boolean.
-        return zerosum_true  # type: ignore
+        return zerosum_true
 
     def eq(self, other: object, static: bool = False) -> bool:
         if not isinstance(other, _ZeroSum):

numpyro/distributions/distribution.py

@@ -426,7 +426,7 @@ def variance(self) -> Array:
     def mode(self) -> Array:
         raise NotImplementedError
 
-    def _validate_sample(self, value: ArrayLike) -> Array:
+    def _validate_sample(self, value: ArrayLike) -> ArrayLike:
         assert self.support is not None
         mask = self.support(value)
         if not_jax_tracer(mask):
@@ -909,7 +909,7 @@ def log_prob(self, value: ArrayLike) -> Array:
         batch_shape = lax.broadcast_shapes(batch_shape, self.batch_shape)
         return jnp.zeros(batch_shape)
 
-    def _validate_sample(self, value: ArrayLike) -> Array:
+    def _validate_sample(self, value: ArrayLike) -> ArrayLike:
         mask = super(ImproperUniform, self)._validate_sample(value)
         batch_dim = jnp.ndim(value) - len(self.event_shape)
         if batch_dim < jnp.ndim(mask):