def call(self, inputs):
out = self.dense(inputs)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>
Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].

        Returns:
          A batch of MultivariateNormalDiag distributions with event shape
          [latent_size], batch shape [..., batch_size, timesteps], and
          sample shape [sample_shape, ..., batch_size, timesteps,
          latent_size].