Remove explicit type casting in quantized.

PiperOrigin-RevId: 462610652
google-deepmind · Jul 22, 2022 · f3a1b88 · f3a1b88
1 parent 0ecad05
commit f3a1b88
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/distrax/_src/distributions/quantized.py b/distrax/_src/distributions/quantized.py
@@ -14,7 +14,7 @@
 # ==============================================================================
 """Quantized distribution."""
 
-from typing import cast, Optional, Tuple
+from typing import Optional, Tuple
 
 import chex
 from distrax._src.distributions import distribution as base_distribution
@@ -107,9 +107,7 @@ def high(self) -> Optional[Array]:
   @property
   def event_shape(self) -> Tuple[int, ...]:
     """Shape of event of distribution samples."""
-    event_shape = self.distribution.event_shape
-    # TODO(b/149413467): Remove explicit casting when resolved.
-    return cast(Tuple[int, ...], event_shape)
+    return self.distribution.event_shape
 
   @property
   def batch_shape(self) -> Tuple[int, ...]: