Passing tf.RaggedTensor to tfp.Distribution's methods in Python Tensorflow

Suppose I have a tensorflow-probability distribution

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

distr = tfd.Normal(loc=0, scale=1)

and I want to evaluate negative log-likelihood on a ragged tensor

>>> type(rt)
tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor
>>> rt.shape
TensorShape([50, None])
>>> distr.log_prob(rt)  # Expected shape: [50, None]
ValueError: TypeError: object of type 'RaggedTensor' has no len()

I would like to use it in a loss function during model training (which means that I can't unstack the tensor's batch dimension):

>>> def negloglike(y, distr):
        mean_over_samples = tf.reduce_mean(distr.log_prob(y), axis=-1)  # shape [batch_size]
        return -tf.reduce_mean(mean_over_samples)

Attempt to do this produces an error

<...>
line 2, in negloglike  *
        mean_over_samples = tf.reduce_mean(model.log_prob(y))
<...>
TypeError: Failed to convert elements of tf.RaggedTensor <...> to Tensor. <..>

I tried padding RaggedTensor with NaNs to get a regular tf.Tensor, followed with masking NaN values after log_prob(y), but it made model weights go NaN as well.

Is there any way to overcome this obstacle?

I was able to find the solution that at least works. To keep it simple, the most important detail is tf.RaggedTensor.flat_values attribute. It transformed the RaggedTensor to regular Tensor by concatenating all the rows. (It works in a more complicated manner for RaggedTensors with multiple ragged dimensions.)

Another thing is that I added another dimension to the distribution's batch_shape to make sure that distribution.log_prob() accepts flattened RaggedTensor. However, it increases the number of computed probability values by approximately batch_size. Most of them are ditched on the next step.

Finally, to obtain 1D Tensor again, I used tf.gather_nd. It picked correct values from the 2D log_probs_2d to 1D log_probs. The only operation left is averaging.

Full code is as follows:

def negloglike(y, distr):
    '''
    y: tf.RaggedTensor, shape (batch_size, None)
    distr: batch_shape==(None, 1), event_shape==(,)
    '''

    batch_size = tf.shape(y)[0]
    segment_ids = tf.repeat(tf.range(batch_size), y.row_lengths())  # shape (total_samples, )
    total_samples = tf.shape(segment_ids)[0]
    
    flattened = y.flat_values  # shape (total_samples, )
    log_probs_2d = distr.prob(flattened) # shape (batch_size, total_samples, )
    ### calculating too many unnecessary values (\times batch_size more than needed)

    log_probs = tf.gather_nd(indices=tf.stack([segment_ids,
                                               tf.range(total_samples)], axis=1), 
                             params=log_probs_2d)  # shape (total_samples, )
    
    mean_over_samples = tf.math.segment_mean(log_probs, segment_ids)  # shape (batch_size, )
    
    return -tf.reduce_mean(mean_over_samples)