From 26cf56cddd7ba8123794a6c1fe2de2b31f9a1b30 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Fri, 11 Aug 2023 12:13:53 -0700 Subject: [PATCH] [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 | 4 ++-- distrax/_src/distributions/categorical.py | 4 ++-- distrax/_src/distributions/epsilon_greedy.py | 4 +++- distrax/_src/distributions/greedy.py | 4 +++- distrax/_src/distributions/multinomial.py | 6 +++--- distrax/_src/distributions/one_hot_categorical.py | 5 +++-- distrax/_src/distributions/softmax.py | 4 +++- distrax/_src/distributions/transformed.py | 1 + 8 files changed, 20 insertions(+), 12 deletions(-) diff --git a/distrax/_src/distributions/bernoulli.py b/distrax/_src/distributions/bernoulli.py index 662e764..6806346 100644 --- a/distrax/_src/distributions/bernoulli.py +++ b/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: jnp.dtype | type[Any] = int): """Initializes a Bernoulli distribution. Args: diff --git a/distrax/_src/distributions/categorical.py b/distrax/_src/distributions/categorical.py index 2ed8be4..55416b2 100644 --- a/distrax/_src/distributions/categorical.py +++ b/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: jnp.dtype | type[Any] = int): """Initializes a Categorical distribution. Args: diff --git a/distrax/_src/distributions/epsilon_greedy.py b/distrax/_src/distributions/epsilon_greedy.py index 14b077b..d738bac 100644 --- a/distrax/_src/distributions/epsilon_greedy.py +++ b/distrax/_src/distributions/epsilon_greedy.py @@ -14,6 +14,8 @@ # ============================================================================== """Epsilon-Greedy distributions with respect to a set of preferences.""" +from typing import Any + 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: jnp.dtype | type[Any] = int): """Initializes an EpsilonGreedy distribution. Args: diff --git a/distrax/_src/distributions/greedy.py b/distrax/_src/distributions/greedy.py index e5a8873..0bff51d 100644 --- a/distrax/_src/distributions/greedy.py +++ b/distrax/_src/distributions/greedy.py @@ -14,6 +14,8 @@ # ============================================================================== """Greedy distributions with respect to a set of preferences.""" +from typing import Any + import chex from distrax._src.distributions import categorical from distrax._src.distributions import distribution @@ -37,7 +39,7 @@ 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: jnp.dtype | type[Any] = int): """Initializes a Greedy distribution. Args: diff --git a/distrax/_src/distributions/multinomial.py b/distrax/_src/distributions/multinomial.py index 3df2908..08ea0c6 100644 --- a/distrax/_src/distributions/multinomial.py +++ b/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: 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: int | Array, n: int, logits: Array, dtype: jnp.dtype) -> Array: """Sample method for a Multinomial with integer `total_count`.""" diff --git a/distrax/_src/distributions/one_hot_categorical.py b/distrax/_src/distributions/one_hot_categorical.py index 75004e5..f30bcf3 100644 --- a/distrax/_src/distributions/one_hot_categorical.py +++ b/distrax/_src/distributions/one_hot_categorical.py @@ -14,7 +14,7 @@ # ============================================================================== """OneHotCategorical distribution.""" -from typing import Optional, Tuple +from typing import Any, Optional, Tuple 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: 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) diff --git a/distrax/_src/distributions/softmax.py b/distrax/_src/distributions/softmax.py index 86ded3f..5fc3b6b 100644 --- a/distrax/_src/distributions/softmax.py +++ b/distrax/_src/distributions/softmax.py @@ -14,6 +14,8 @@ # ============================================================================== """Softmax distribution.""" +from typing import Any + 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: jnp.dtype | type[Any] = int): """Initializes a Softmax distribution. Args: diff --git a/distrax/_src/distributions/transformed.py b/distrax/_src/distributions/transformed.py index f066a28..538b59a 100644 --- a/distrax/_src/distributions/transformed.py +++ b/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