Let DeferredTensor know the static shape of the broadcasted inflated_loc.

PiperOrigin-RevId: 475879856

Author: ThomasColthurst

tensorflow_probability/python/distributions/inflated.py

@@ -148,7 +148,8 @@ def __init__(self,
               probs_or_logits,
               # pylint: disable=g-long-lambda
               lambda _: tf.broadcast_to(self._inflated_loc,
-                                        ps.shape(probs_or_logits))),
+                                        ps.shape(probs_or_logits)),
+              shape=probs_or_logits.shape),
           atol=self._inflated_loc_atol,
           rtol=self._inflated_loc_rtol,
           validate_args=validate_args,