def call(self, inputs):
    """
    """
    # TODO(dusenberrymw): Remove these reshaping commands after b/113126249 is
    # fixed.
    collapsed_shape = tf.concat(([-1], tf.shape(input=inputs)[-2:]), axis=0)
    out = tf.reshape(inputs, collapsed_shape)  # (sample*batch_size, T, hidden)
    out = self.bilstm(out)  # (sample*batch_size, hidden)
    expanded_shape = tf.concat((tf.shape(input=inputs)[:-2], [-1]), axis=0)
    out = tf.reshape(out, expanded_shape)  # (sample, batch_size, hidden)
    out = self.output_layer(out)  # (sample, batch_size, 2*latent_size)
    loc = out[..., :self.latent_size]
    scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)