def METHOD_NAME(self, *args, expected_batch_shape=None, **kwargs):
    expected_batch_shape = expected_batch_shape or self.batch_shape
    expected_batch_shape = expected_batch_shape[:-1]
    super().METHOD_NAME(*args, expected_batch_shape=expected_batch_shape, **kwargs)