def __init__(self, dimensions, hidden_size):
super(LearnableMultivariateNormalDiagCell, self).__init__()<EOL>self.dimensions = dimensions<EOL>self.hidden_size = hidden_size<EOL>self.lstm_cell = tf.keras.layers.LSTMCell(hidden_size)<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*dimensions)<EOL>
Constructs a learnable multivariate diagonal normal cell.

        Args:
          dimensions: An integer corresponding to the dimensionality of the
            distribution.
          hidden_size: Dimensionality of the LSTM function parameters.