code: def do_elbo_test(
      self,
      repa1,
      repa2,
      n_steps,
      prec,
      lr,
      use_nn_baseline,
      use_decaying_avg_baseline,
  ):
      logger.info(
          " - - - - - DO NORMALNORMALNORMAL ELBO TEST - - - - - -"
      )
      logger.info(
          "[reparameterized = %s, %s; nn_baseline = %s, decaying_baseline = %s]"
          % (repa1, repa2, use_nn_baseline, use_decaying_avg_baseline)
      )
      pyro.clear_param_store()
      Normal1 = dist.Normal if repa1 else fakes.NonreparameterizedNormal
      Normal2 = dist.Normal if repa2 else fakes.NonreparameterizedNormal
  
      if use_nn_baseline:
  
          class VanillaBaselineNN(nn.Module):
              def __init__(self, dim_input, dim_h):
                  super().__init__()
                  self.lin1 = nn.Linear(dim_input, dim_h)
                  self.lin2 = nn.Linear(dim_h, 2)
                  self.sigmoid = nn.Sigmoid()
  
              def forward(self, x):
                  h = self.sigmoid(self.lin1(x))
                  return self.lin2(h)
  
          loc_prime_baseline = pyro.module(
              "loc_prime_baseline", VanillaBaselineNN(2, 5)
          )
      else:
          loc_prime_baseline = None
  
      def model():
          with pyro.plate("plate", 2):
              loc_latent_prime = pyro.sample(
                  "loc_latent_prime",
                  Normal1(self.loc0, torch.pow(self.lam0, -0.5)),
              )
              loc_latent = pyro.sample(
                  "loc_latent",
                  Normal2(loc_latent_prime, torch.pow(self.lam0, -0.5)),
              )
              with pyro.plate("data", len(self.data)):
                  pyro.sample(
                      "obs",
                      dist.Normal(
                          loc_latent, torch.pow(self.lam, -0.5)
                      ).expand_by(self.data.shape[:1]),
                      obs=self.data,
                  )
          return loc_latent
  
      # note that the exact posterior is not mean field!
      def guide():
          loc_q = pyro.param(
              "loc_q", self.analytic_loc_n.expand(2) + 0.334
          )
          log_sig_q = pyro.param(
              "log_sig_q", self.analytic_log_sig_n.expand(2) - 0.29
          )
          loc_q_prime = pyro.param(
              "loc_q_prime", torch.tensor([-0.34, 0.52])
          )
          kappa_q = pyro.param("kappa_q", torch.tensor([0.74]))
          log_sig_q_prime = pyro.param(
              "log_sig_q_prime", -0.5 * torch.log(1.2 * self.lam0)
          )
          sig_q, sig_q_prime = torch.exp(log_sig_q), torch.exp(
              log_sig_q_prime
          )
          with pyro.plate("plate", 2):
              loc_latent = pyro.sample(
                  "loc_latent",
                  Normal2(loc_q, sig_q),
                  infer=dict(
                      baseline=dict(
                          use_decaying_avg_baseline=use_decaying_avg_baseline
                      )
                  ),
              )
              pyro.sample(
                  "loc_latent_prime",
                  Normal1(
                      kappa_q.expand_as(loc_latent) * loc_latent
                      + loc_q_prime,
                      sig_q_prime,
                  ),
                  infer=dict(
                      baseline=dict(
                          nn_baseline=loc_prime_baseline,
                          nn_baseline_input=loc_latent,
                          use_decaying_avg_baseline=use_decaying_avg_baseline,
                      )
                  ),
              )
              with pyro.plate("data", len(self.data)):
                  pass
  
          return loc_latent
  
      adam = optim.Adam({"lr": 0.0015, "betas": (0.97, 0.999)})
      svi = SVI(model, guide, adam, loss=TraceGraph_ELBO())
  
      for k in range(n_steps):
          svi.step()
  
          loc_error = param_mse("loc_q", self.analytic_loc_n)
          log_sig_error = param_mse(
              "log_sig_q", self.analytic_log_sig_n
          )
          loc_prime_error = param_mse("loc_q_prime", 0.5 * self.loc0)
          kappa_error = param_mse("kappa_q", 0.5 * torch.ones(1))
          log_sig_prime_error = param_mse(
              "log_sig_q_prime", -0.5 * torch.log(2.0 * self.lam0)
          )
  
          if k % 500 == 0:
              logger.debug(
                  "errors:  %.4f, %.4f" % (loc_error, log_sig_error)
              )
              logger.debug(
                  ", %.4f, %.4f"
                  % (loc_prime_error, log_sig_prime_error)
              )
              logger.debug(", %.4f" % kappa_error)
  
      assert_equal(0.0, loc_error, prec=prec)
      assert_equal(0.0, log_sig_error, prec=prec)
      assert_equal(0.0, loc_prime_error, prec=prec)
      assert_equal(0.0, log_sig_prime_error, prec=prec)
      assert_equal(0.0, kappa_error, prec=prec)
  
imports: import logging
  from unittest import TestCase
  import numpy
  import pytest
  import torch
  from torch import nn
  import pyro
  import pyro.distributions
  import pyro.optim
  from pyro.distributions.testing import fakes
  from pyro.infer import SVI, TraceGraph_ELBO
  from tests.common import assert_equal
label_desc: Use image node_id 6 for calling the NormalNormalNormalTests obj's underlying member method code with example usage: obj.do_elbo_test(repa1, repa2, n_steps, prec, lr, use_nn_baseline, use_decaying_avg_baseline) and returns: loc_latent, loc_latent, self