From e131ec22593b83e011153ffb1290189b92e0e18e Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Mon, 14 Sep 2026 15:09:17 -0700 Subject: [PATCH] Suppress pyrefly errors related to JAX scalar types PiperOrigin-RevId: 981385491 --- distrax/_src/distributions/multinomial.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/distrax/_src/distributions/multinomial.py b/distrax/_src/distributions/multinomial.py index 3868818..16ee039 100644 --- a/distrax/_src/distributions/multinomial.py +++ b/distrax/_src/distributions/multinomial.py @@ -197,7 +197,7 @@ def entropy(self) -> Array: @staticmethod def _entropy_scalar( total_count: int, probs: Array, log_of_probs: Array - ) -> Union[jnp.float32, jnp.float64]: + ) -> Union[jnp.float32, jnp.float64]: # pyrefly: ignore[not-a-type] """Calculates the entropy for a Multinomial with integer `total_count`.""" # Constant factors in the entropy. xi = jnp.arange(total_count + 1, dtype=probs.dtype) @@ -230,7 +230,7 @@ def _entropy_scalar( @staticmethod def _entropy_scalar_with_lax( total_count: int, probs: Array, log_of_probs: Array - ) -> Union[jnp.float32, jnp.float64]: + ) -> Union[jnp.float32, jnp.float64]: # pyrefly: ignore[not-a-type] """Like `_entropy_scalar`, but uses a lax while loop.""" dtype = probs.dtype