def call(self, inputs, state):
    """
    """
    # In order to allow the user to pass in a single example without a batch
    # dimension, we always expand the input to at least two dimensions, then
    # fix the output shape to remove the batch dimension if necessary.
    original_shape = inputs.shape
    if len(original_shape) < 2:
      inputs = tf.reshape(inputs, [1, -1])
    out, state = self.lstm_cell(inputs, state)
    out = self.output_layer(out)
    correct_shape = tf.concat((original_shape[:-1], tf.shape(input=out)[-1:]),
                              0)
    out = tf.reshape(out, correct_shape)
    loc = out[..., :self.dimensions]
    scale_diag = tf.nn.softplus(out[..., self.dimensions:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag), state