class_docstr: Multivariate diagonal normal distribution RNN cell.
  
  The model is an LSTM-based recurrent function that computes the
  parameters for a multivariate normal distribution at each timestep
  `t`.
class_docstr_tok: ["Multivariate","diagonal","normal","distribution","RNN","cell",".","The","model","is","an","LSTM-based","recurrent","function","that","computes","the","parameters","for","a","multivariate","normal","distribution","at","each","timestep","`","t","`","."]
docstr: Runs the model to generate a distribution for a single timestep.
  
  This generates a batched MultivariateNormalDiag distribution using
  the output of the recurrent model at the current timestep to
  parameterize the distribution.
  
  Args:
    inputs: The sampled value of `z` at the previous timestep, i.e.,
      `z_{t-1}`, of shape [..., dimensions].
      `z_0` should be set to the empty matrix.
    state: A tuple containing the (hidden, cell) state.
  
  Returns:
    A tuple of a MultivariateNormalDiag distribution, and the state of
    the recurrent function at the end of the current timestep. The
    distribution will have event shape [dimensions], batch shape
    [...], and sample shape [sample_shape, ..., dimensions].
code: def call(self, inputs, state):
      """Runs the model to generate a distribution for a single timestep.
  
      This generates a batched MultivariateNormalDiag distribution using
      the output of the recurrent model at the current timestep to
      parameterize the distribution.
  
      Args:
        inputs: The sampled value of `z` at the previous timestep, i.e.,
          `z_{t-1}`, of shape [..., dimensions].
          `z_0` should be set to the empty matrix.
        state: A tuple containing the (hidden, cell) state.
  
      Returns:
        A tuple of a MultivariateNormalDiag distribution, and the state of
        the recurrent function at the end of the current timestep. The
        distribution will have event shape [dimensions], batch shape
        [...], and sample shape [sample_shape, ..., dimensions].
      """
      # 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 = tf1.reshape(inputs, [1, -1])
      out, state = self.lstm_cell(inputs, state)
      out = self.output_layer(out)
      # correct_shape = tf1.concat(
      #     (original_shape[:-1], tf1.shape(input=out)[-1:]), 0)
      # out = tf1.reshape(out, correct_shape)
      loc = out[..., : self.dimensions]
      scale_diag = (
          tf1.nn.softplus(out[..., self.dimensions :]) + 1e-5
      )  # keep > 0
      return (
          tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag),
          state,
      )
  
imports: from datetime import datetime
  import functools
  from absl import app
  from absl import flags
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.examples import sprites_dataset
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 3 for calling the LearnableMultivariateNormalDiagCell obj's underlying member method code with example usage: obj.call(inputs, state) and returns: tfd,  state