Closes #297

Description of the problem:

When not using auto_batching, the event_shape of the distributions and potentials must match. The reason is, when the log-likelihood is calculated, it is not possible to reduce the sum over all dimensions, as the batch_shape has to be preserved. This works well for distributions, however, transformations add a potential as log_det_jacobian with too high dimensionality. It is assumed in SamplingState.collect_unreduced_log_prob that each potential and distribution consist only of batch_dim as shape, however the backward_log_det_jacobian also includes the event_dim.

Solution

An event_ndims attribute is added to the transform class that saves the number of event dimensions of the distributions it belongs to. This number is passed to the log_det_jacobian function implemented in tensorflow_probability, which then reduces over these axes. At the end of the initialization of a pymc4 distribution, the event_ndims attribute of the transform is set by reading out the event_shape attribute of the just initialized tfp.distributions.distribution.
There is some additional code to take care of the case when the transformation changes the number of dimensions. This change of dimensions is defined in the inverse_event_ndims method of the tfp.bijectors.bijector, which is used to calculate the number of dimensions of the transformed variable.

Comparison to tensorflow_probability

In tensorflow_probability, in TransformedTransitionKernel, a similar problem has been already been solved. Their solution was to calculate the rank of the untransformed likelihood function, which is equal to the length of the batch_shape and subtract it from the rank of the input to each bijector to get event_ndims link. This is however done at the sampling stage. I suppose that theoretically, one could also implement something similar in PyMC4 in the SamplingExecutor, however, I think that it is cleaner to do it already at the initialization stage because the logic for the event_shapes of the distributions is also implemented there.

Implementation details

• The event_ndims is read out from the tfp.distribution with prefer_static, which is copied from link.
• One can provide a event_ndims as initialization parameter to a transform object, which than hinders the automatic setting of event_ndims by the distribution class.

Other Changes

• Fix requirements of xarray <= 0.16.0 because of xarray 0.16.1 breaks compatibility arviz-devs/arviz#1387
• Fixes some doctests such that the tests run through without errors