Fix a couple of glitches with bijector batch shape.

1. ScaleMatvec* bijectors were missing annotations because I'd accidentally added them to the Block variant rather than to the base class.
2. Passing `np.int32` or `np.int64` values for event_ndims could lead to dtype mismatch errors.

PiperOrigin-RevId: 379987638

Author: davmre

tensorflow_probability/python/bijectors/BUILD

@@ -254,6 +254,7 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensor_util",
         "//tensorflow_probability/python/internal:tensorshape_util",

tensorflow_probability/python/bijectors/bijector_test.py

@@ -371,6 +371,8 @@ class BijectorBatchShapesTest(test_util.TestCase):
        lambda: tfb.Scale(tf.ones([4, 2])), None),
       ('sigmoid',
        lambda: tfb.Sigmoid(low=tf.zeros([3]), high=tf.ones([4, 1])), None),
+      ('scale_matvec',
+       lambda: tfb.ScaleMatvecDiag([[0.], [3.]]), None),
       ('invert',
        lambda: tfb.Invert(tfb.ScaleMatvecDiag(tf.ones([2, 1]))), None),
       ('reshape',
@@ -416,6 +418,14 @@ def test_batch_shape_matches_output_shapes(self,
     self.assertAllEqual(batch_shape_tensor_x, batch_shape_tensor_y)
     self.assertAllEqual(batch_shape_tensor_x, batch_shape_x)
 
+    # Check that we're robust to integer type.
+    batch_shape_tensor_x64 = bijector.experimental_batch_shape_tensor(
+        x_event_ndims=tf.nest.map_structure(np.int64, x_event_ndims))
+    batch_shape_tensor_y64 = bijector.experimental_batch_shape_tensor(
+        y_event_ndims=tf.nest.map_structure(np.int64, y_event_ndims))
+    self.assertAllEqual(batch_shape_tensor_x64, batch_shape_tensor_y64)
+    self.assertAllEqual(batch_shape_tensor_x64, batch_shape_x)
+
     # Pushing a value through the bijector should return a Tensor(s) with
     # the expected batch shape...
     xs = tf.nest.map_structure(lambda nd: tf.ones([1] * nd), x_event_ndims)

tensorflow_probability/python/bijectors/scale_matvec_linear_operator.py

@@ -22,6 +22,7 @@
 
 from tensorflow_probability.python.bijectors import bijector
 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 tensor_util
 from tensorflow_probability.python.internal import tensorshape_util
@@ -46,6 +47,10 @@ def adjoint(self):
     """`bool` indicating whether this class uses `self.scale` or its adjoint."""
     return self._adjoint
 
+  @classmethod
+  def _parameter_properties(cls, dtype):
+    return dict(scale=parameter_properties.BatchedComponentProperties())
+
   def _forward(self, x):
     return self.scale.matvec(x, adjoint=self.adjoint)
 
@@ -237,10 +242,6 @@ def _inverse_event_shape_tensor(self, output_shape):
       return _cumulative_broadcast_dynamic(output_shape)
     return output_shape
 
-  @classmethod
-  def _parameter_properties(cls, dtype):
-    return {}
-
 
 def _cumulative_broadcast_static(event_shape):
   broadcast_shapes = [s[:-1] for s in event_shape]

tensorflow_probability/python/internal/batch_shape_lib.py

@@ -137,6 +137,7 @@ def get_base_shape(x):
 
 def slice_batch_shape_tensor(base_shape, event_ndims):
   base_shape = ps.convert_to_shape_tensor(base_shape, dtype_hint=np.int32)
+  event_ndims = ps.convert_to_shape_tensor(event_ndims, dtype_hint=np.int32)
   base_rank = ps.rank_from_shape(base_shape)
   return base_shape[:(base_rank -
                       # Don't try to slice away more ndims than the parameter