Use prefer_static.shape over tf.shape in inflated.py's broadcasting of

inflated_loc, to be support jax.jit.

PiperOrigin-RevId: 475822622

Author: ThomasColthurst

tensorflow_probability/python/distributions/BUILD

@@ -1014,6 +1014,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
+        "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:samplers",
         "//tensorflow_probability/python/internal:tensor_util",
         "//tensorflow_probability/python/util:deferred_tensor",

tensorflow_probability/python/distributions/inflated.py

@@ -28,6 +28,7 @@
 from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
+from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import samplers
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.util.deferred_tensor import DeferredTensor
@@ -147,9 +148,7 @@ def __init__(self,
               probs_or_logits,
               # pylint: disable=g-long-lambda
               lambda _: tf.broadcast_to(self._inflated_loc,
-                                        tf.shape(probs_or_logits))
-              # pylint: enable=g-long-lambda
-          ),
+                                        ps.shape(probs_or_logits))),
           atol=self._inflated_loc_atol,
           rtol=self._inflated_loc_rtol,
           validate_args=validate_args,