code: def test_elbo_beta_kl(aux_total_size):
      mu0 = 1.5
      sigma = 1.0
      y_obs = np.array([1.6, 1.4])
      beta = len(y_obs) / float(aux_total_size)
  
      with pytensor.config.change_flags(
          floatX="float64", warn_float64="ignore"
      ):
          post_mu = np.array([1.88], dtype=pytensor.config.floatX)
          post_sigma = np.array([1], dtype=pytensor.config.floatX)
  
          with pm.Model():
              mu = pm.Normal("mu", mu=mu0, sigma=sigma)
              pm.Normal(
                  "y",
                  mu=mu,
                  sigma=1,
                  observed=y_obs,
                  total_size=aux_total_size,
              )
              # Create variational gradient tensor
              mean_field_1 = MeanField()
              mean_field_1.scale_cost_to_minibatch = True
              mean_field_1.shared_params["mu"].set_value(post_mu)
              mean_field_1.shared_params["rho"].set_value(
                  np.log(np.exp(post_sigma) - 1)
              )
  
              with pytensor.config.change_flags(
                  compute_test_value="off"
              ):
                  elbo_via_total_size_scaled = -pm.operators.KL(
                      mean_field_1
                  )()(10000)
  
          with pm.Model():
              mu = pm.Normal("mu", mu=mu0, sigma=sigma)
              pm.Normal("y", mu=mu, sigma=1, observed=y_obs)
              # Create variational gradient tensor
              mean_field_3 = MeanField()
              mean_field_3.shared_params["mu"].set_value(post_mu)
              mean_field_3.shared_params["rho"].set_value(
                  np.log(np.exp(post_sigma) - 1)
              )
  
              with pytensor.config.change_flags(
                  compute_test_value="off"
              ):
                  elbo_via_beta_kl = -pm.operators.KL(
                      mean_field_3, beta=beta
                  )()(10000)
  
          np.testing.assert_allclose(
              elbo_via_total_size_scaled.eval(),
              elbo_via_beta_kl.eval(),
              rtol=0,
              atol=1e-1,
          )
  
file_path: pymc/tests/variational/test_approximations.py
imports: import numpy
  import pytensor
  import pytest
  import pymc
  from pymc.variational.approximations import Empirical, MeanField
  from tests import models
label_desc: Use image node_id 5 for calling a global function with example usage: test_elbo_beta_kl(aux_total_size) without return types