[JAX] Fix incorrect type annotations.

An upcoming change to JAX will teach pytype more accurate types for functions in the jax.numpy module. This reveals a number of type errors in downstream users of JAX. In particular, pytype is able to infer `jax.Array` accurately as a type in many more cases.

PiperOrigin-RevId: 556042684

distrax/_src/distributions/bernoulli.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Bernoulli distribution."""
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import chex
 from distrax._src.distributions import distribution
@@ -44,7 +44,7 @@ class Bernoulli(distribution.Distribution):
   def __init__(self,
                logits: Optional[Numeric] = None,
                probs: Optional[Numeric] = None,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes a Bernoulli distribution.
 
     Args:

distrax/_src/distributions/categorical.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Categorical distribution."""
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import chex
 from distrax._src.distributions import distribution
@@ -38,7 +38,7 @@ class Categorical(distribution.Distribution):
   def __init__(self,
                logits: Optional[Array] = None,
                probs: Optional[Array] = None,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes a Categorical distribution.
 
     Args:

distrax/_src/distributions/epsilon_greedy.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Epsilon-Greedy distributions with respect to a set of preferences."""
 
+from typing import Any, Union
+
 import chex
 from distrax._src.distributions import categorical
 from distrax._src.distributions import distribution
@@ -46,7 +48,7 @@ class EpsilonGreedy(categorical.Categorical):
   def __init__(self,
                preferences: Array,
                epsilon: float,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes an EpsilonGreedy distribution.
 
     Args:

distrax/_src/distributions/greedy.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Greedy distributions with respect to a set of preferences."""
 
+from typing import Any, Union
+
 import chex
 from distrax._src.distributions import categorical
 from distrax._src.distributions import distribution
@@ -37,7 +39,9 @@ class Greedy(categorical.Categorical):
   all other indices will be assigned a probability of zero.
   """
 
-  def __init__(self, preferences: Array, dtype: jnp.dtype = int):
+  def __init__(
+      self, preferences: Array, dtype: Union[jnp.dtype, type[Any]] = int
+  ):
     """Initializes a Greedy distribution.
 
     Args:

distrax/_src/distributions/multinomial.py

@@ -17,7 +17,7 @@
 import functools
 import operator
 
-from typing import Tuple, Optional, Union
+from typing import Any, Tuple, Optional, Union
 
 import chex
 from distrax._src.distributions import distribution
@@ -46,7 +46,7 @@ def __init__(self,
                total_count: Numeric,
                logits: Optional[Array] = None,
                probs: Optional[Array] = None,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes a Multinomial distribution.
 
     Args:
@@ -162,7 +162,7 @@ def _sample_n(self, key: PRNGKey, n: int) -> Array:
 
   @staticmethod
   def _sample_n_scalar(
-      key: PRNGKey, total_count: int, n: int, logits: Array,
+      key: PRNGKey, total_count: Union[int, Array], n: int, logits: Array,
       dtype: jnp.dtype) -> Array:
     """Sample method for a Multinomial with integer `total_count`."""
 

distrax/_src/distributions/one_hot_categorical.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """OneHotCategorical distribution."""
 
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple, Union
 
 import chex
 from distrax._src.distributions import categorical
@@ -40,7 +40,7 @@ class OneHotCategorical(categorical.Categorical):
   def __init__(self,
                logits: Optional[Array] = None,
                probs: Optional[Array] = None,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes a OneHotCategorical distribution.
 
     Args:
@@ -80,6 +80,7 @@ def prob(self, value: EventT) -> Array:
   def mode(self) -> Array:
     """Calculates the mode."""
     preferences = self._probs if self._logits is None else self._logits
+    assert preferences is not None
     greedy_index = jnp.argmax(preferences, axis=-1)
     return jax.nn.one_hot(greedy_index, self.num_categories).astype(self._dtype)
 

distrax/_src/distributions/softmax.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Softmax distribution."""
 
+from typing import Any, Union
+
 import chex
 from distrax._src.distributions import categorical
 from distrax._src.distributions import distribution
@@ -35,7 +37,7 @@ class Softmax(categorical.Categorical):
   def __init__(self,
                logits: Array,
                temperature: float = 1.,
-               dtype: jnp.dtype = int):
+               dtype: Union[jnp.dtype, type[Any]] = int):
     """Initializes a Softmax distribution.
 
     Args:

distrax/_src/distributions/transformed.py

@@ -139,6 +139,7 @@ def dtype(self) -> jnp.dtype:
     """See `Distribution.dtype`."""
     if self._dtype is None:
       self._infer_shapes_and_dtype()
+    assert self._dtype is not None  # By _infer_shapes_and_dtype()
     return self._dtype
 
   @property