Model.__call__ with VariationalGaussianProcess fails when input length is 1

The simplest way to reproduce is to take https://github.com/tensorflow/probability/blob/r0.9/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_Layers_Regression.ipynb , change it to install TF 2.1 (otherwise TFP won't run due to dependency TF>2.0).

Then run this as the last cell model(x_tst[:1])
Result:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-53-44d7272f4074> in <module>()
----> 1 model(x_tst[:1])

20 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    862           with base_layer_utils.autocast_context_manager(
    863               self._compute_dtype):
--> 864             outputs = self.call(cast_inputs, *args, **kwargs)
    865           self._handle_activity_regularization(inputs, outputs)
    866           self._set_mask_metadata(inputs, outputs, input_masks)

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/sequential.py in call(self, inputs, training, mask)
    271       if not self.built:
    272         self._init_graph_network(self.inputs, self.outputs, name=self.name)
--> 273       return super(Sequential, self).call(inputs, training=training, mask=mask)
    274 
    275     outputs = inputs  # handle the corner case where self.layers is empty

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/network.py in call(self, inputs, training, mask)
    713     return self._run_internal_graph(
    714         inputs, training=training, mask=mask,
--> 715         convert_kwargs_to_constants=base_layer_utils.call_context().saving)
    716 
    717   def compute_output_shape(self, input_shape):

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/network.py in _run_internal_graph(self, inputs, training, mask, convert_kwargs_to_constants)
    890 
    891           # Compute outputs.
--> 892           output_tensors = layer(computed_tensors, **kwargs)
    893 
    894           # Update tensor_dict.

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py in __call__(self, inputs, *args, **kwargs)
    243     self._enter_dunder_call = True
    244     distribution, _ = super(DistributionLambda, self).__call__(
--> 245         inputs, *args, **kwargs)
    246     self._enter_dunder_call = False
    247     return distribution

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    862           with base_layer_utils.autocast_context_manager(
    863               self._compute_dtype):
--> 864             outputs = self.call(cast_inputs, *args, **kwargs)
    865           self._handle_activity_regularization(inputs, outputs)
    866           self._set_mask_metadata(inputs, outputs, input_masks)

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py in call(self, inputs, *args, **kwargs)
    249   def call(self, inputs, *args, **kwargs):
    250     distribution, value = super(DistributionLambda, self).call(
--> 251         inputs, *args, **kwargs)
    252     # We always save the most recently built distribution for variable tracking
    253     # purposes.

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/layers/core.py in call(self, inputs, mask, training)
    858     with backprop.GradientTape(watch_accessed_variables=True) as tape,\
    859         variable_scope.variable_creator_scope(_variable_creator):
--> 860       result = self.function(inputs, **kwargs)
    861     self._check_variables(created_variables, tape.watched_variables())
    862     return result

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py in _fn(*fargs, **fkwargs)
    183       # We'd prefer to call ops.convert_to_tensor_or_composite but do not,
    184       # favoring our own non-public API over TF's.
--> 185       value = distribution._value()  # pylint: disable=protected-access
    186 
    187       # TODO(b/126056144): Remove silent handle once we identify how/why Keras

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/internal/distribution_tensor_coercible.py in _value(self, dtype, name, as_ref)
    120       with self._name_and_control_scope('value'):
    121         self._concrete_value = (self._convert_to_tensor_fn(self)
--> 122                                 if callable(self._convert_to_tensor_fn)
    123                                 else self._convert_to_tensor_fn)
    124         if (not tf.is_tensor(self._concrete_value) and

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/distribution.py in sample(self, sample_shape, seed, name, **kwargs)
    856       samples: a `Tensor` with prepended dimensions `sample_shape`.
    857     """
--> 858     return self._call_sample_n(sample_shape, seed, name, **kwargs)
    859 
    860   def _call_log_prob(self, value, name, **kwargs):

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/transformed_distribution.py in _call_sample_n(self, sample_shape, seed, name, **kwargs)
    417       # work, it is imperative that this is the last modification to the
    418       # returned result.
--> 419       y = self.bijector.forward(x, **bijector_kwargs)
    420       y = self._set_sample_static_shape(y, sample_shape)
    421 

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/bijectors/bijector.py in forward(self, x, name, **kwargs)
   1001       NotImplementedError: if `_forward` is not implemented.
   1002     """
-> 1003     return self._call_forward(x, name, **kwargs)
   1004 
   1005   @classmethod

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/bijectors/bijector.py in _call_forward(self, x, name, **kwargs)
    975       if mapping.y is not None:
    976         return mapping.y
--> 977       mapping = mapping.merge(y=self._forward(x, **kwargs))
    978       # It's most important to cache the y->x mapping, because computing
    979       # inverse(forward(y)) may be numerically unstable / lossy. Caching the

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/bijectors/transpose.py in _forward(self, x)
    195 
    196   def _forward(self, x):
--> 197     return self._transpose(x, self.perm)
    198 
    199   def _event_shape(self, shape, static_perm_to_shape):

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/bijectors/transpose.py in _transpose(self, x, perm)
    267 
    268   def _transpose(self, x, perm):
--> 269     perm = self._make_perm(tf.rank(x), perm)
    270     return tf.transpose(a=x, perm=perm)
    271 

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/bijectors/transpose.py in _make_perm(self, x_rank, perm)
    258     dtype = perm.dtype
    259     perm = tf.concat([
--> 260         tf.range(tf.cast(sample_batch_ndims, dtype)),
    261         tf.cast(
    262             sample_batch_ndims + distribution_util.prefer_static_value(perm),

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/math_ops.py in range(start, limit, delta, dtype, name)
   1590     delta = cast(delta, inferred_dtype)
   1591 
-> 1592     return gen_math_ops._range(start, limit, delta, name=name)
   1593 
   1594 

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/gen_math_ops.py in _range(start, limit, delta, name)
   7119         pass  # Add nodes to the TensorFlow graph.
   7120     except _core._NotOkStatusException as e:
-> 7121       _ops.raise_from_not_ok_status(e, name)
   7122   # Add nodes to the TensorFlow graph.
   7123   _, _, _op, _outputs = _op_def_library._apply_op_helper(

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6626   message = e.message + (" name: " + name if name is not None else "")
   6627   # pylint: disable=protected-access
-> 6628   six.raise_from(core._status_to_exception(e.code, message), None)
   6629   # pylint: enable=protected-access
   6630 

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Requires start <= limit when delta > 0: 0/-1 [Op:Range]

model(x_tst[:2]) works.

In our case it was failing in following line:

distributions = [model.model([batch[k].numpy() for k in model.model.input_names])[1] for batch in dataset.batch(32)]

In case dataset length % 32 == 1 we have the same error, otherwise everything is fine.

Run the Probabilistic_Layers_Regression.ipynb example with TensorFlow 2.1 and reproduce the final-cell failure using model(x_tst[:1]). Start by tracing the stack through tensorflow_probability/python/bijectors/transpose.py, especially _make_perm. Done means the same model call works when the input length is 1, while the existing length-2 case remains working.