Implement required batch shape methods for TFP compatible bijectors.

PiperOrigin-RevId: 382357500

distrax/_src/bijectors/tfp_compatible_bijector.py

@@ -86,9 +86,12 @@ def name(self) -> str:
       """The name of the wrapped bijector."""
       return name_ or f"TFPCompatible{base_bijector.name}"
 
-    @property
-    def experimental_batch_shape(self) -> None:
-      return None
+    def experimental_batch_shape(self, x_event_ndims=None, y_event_ndims=None):
+      raise NotImplementedError()
+
+    def experimental_batch_shape_tensor(
+        self, x_event_ndims=None, y_event_ndims=None):
+      raise NotImplementedError()
 
     def forward_dtype(self, _: jnp.dtype) -> None:
       """Returns None, making no promise regarding dtypes."""

distrax/_src/distributions/multinomial.py

@@ -23,7 +23,7 @@
 from distrax._src.distributions import distribution
 from distrax._src.utils import math
 import jax
-import jax.lax as lax
+from jax import lax
 import jax.numpy as jnp
 from tensorflow_probability.substrates import jax as tfp
 
@@ -75,7 +75,7 @@ def __init__(self,
     else:
       assert self._logits is not None
       probs_batch_shape = self._logits.shape[:-1]
-    self._batch_shape = jax.lax.broadcast_shapes(
+    self._batch_shape = lax.broadcast_shapes(
         probs_batch_shape, self._total_count.shape)
 
   @property