Remove explicit type casting in quantized.

PiperOrigin-RevId: 462610652

Author: hamzamerzic

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
@@ -61,8 +61,7 @@ def __init__(self,
         `distribution` and must not result in additional batch dimensions after
         broadcasting.
     """
-    self._dist: base_distribution.Distribution[Array, Tuple[
-        int, ...], jnp.dtype] = conversion.as_distribution(distribution)
+    self._dist = conversion.as_distribution(distribution)
     if self._dist.event_shape:
       raise ValueError(f'The base distribution must be univariate, but its '
                        f'`event_shape` is {self._dist.event_shape}.')
@@ -107,9 +106,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, ...]: