# Path: epi_forecast_stat_mech/statistical_models/network_models.py
# class PerceptronModule(nn.Module):
# class PlainLinearModule(
#     PerceptronModule.partial(
#         hidden_layer_sizes=(), activation=None, batch_norm_style='none')):
# class LinearModule(
#     PerceptronModule.partial(
#         hidden_layer_sizes=(), activation=None, batch_norm_style='once')):
# class NormalDistributionModel(base.StatisticalModel):
#   def apply(
#       self,
#       inputs,
#       output_size,
#       hidden_layer_sizes=(64, 64),
#       activation=jax.nn.relu,
#       batch_norm_style='once'
#   ):
#   def _output_size_and_unpack_fn(self, output_structure):
#   def init_parameters(self, rng, inputs, output_structure):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, observations):
#   def get_loc_scale(self, parameters, covariates, observations):
#   def predict(self, parameters, covariates, observations):
#   def linear_coefficients(self, parameters):
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


from absl.testing import absltest
from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import network_models
from epi_forecast_stat_mech.statistical_models import tree_util
from jax.config import config
import jax
import jax.numpy as jnp
    """Tests that `log_likelihood` and `predict` return expected shapes."""
    rng = jax.random.PRNGKey(42)
    c_rng, o_rng, init_rng = jax.random.split(rng, 3)
    covariates = jax.random.uniform(c_rng, (num_locations, num_covariates))
    observations = jax.random.uniform(o_rng, (num_locations, num_observations))
    if dict_keys is not None:
      observations = {key: observations for key in dict_keys}

    if error_model == 'full':
      scale_eps = 1E-2
    else:
      # The length of this needs to agree with the flattened number of
      # "observations" per row. I.e. flattened observations is (7, 6)
      # for the 'plugin' test because the two dict keys double the
      # three observations per location.
      scale_eps = jnp.asarray([.1, .2, .3, .4, .5, .6])
    stat_model = network_models.NormalDistributionModel(
        predict_module=predict_module,
        error_model=error_model,
        scale_eps=scale_eps)
    params = stat_model.init_parameters(init_rng, covariates, observations)

    log_likelihood = stat_model.log_likelihood(params, covariates, observations)

    if error_model == 'full':
      actual_shapes = tree_util.tree_map(lambda x: x.shape, log_likelihood)
      expected_shapes = tree_util.tree_map(lambda x: x.shape, observations)
      self.assertEqual(actual_shapes, expected_shapes)
      distributions_over_observations = stat_model.predict(
          params, covariates, observations)

# Path: epi_forecast_stat_mech/statistical_models/network_models.py
# class PerceptronModule(nn.Module):
# class PlainLinearModule(
#     PerceptronModule.partial(
#         hidden_layer_sizes=(), activation=None, batch_norm_style='none')):
# class LinearModule(
#     PerceptronModule.partial(
#         hidden_layer_sizes=(), activation=None, batch_norm_style='once')):
# class NormalDistributionModel(base.StatisticalModel):
#   def apply(
#       self,
#       inputs,
#       output_size,
#       hidden_layer_sizes=(64, 64),
#       activation=jax.nn.relu,
#       batch_norm_style='once'
#   ):
#   def _output_size_and_unpack_fn(self, output_structure):
#   def init_parameters(self, rng, inputs, output_structure):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, observations):
#   def get_loc_scale(self, parameters, covariates, observations):
#   def predict(self, parameters, covariates, observations):
#   def linear_coefficients(self, parameters):
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


# Path: epi_forecast_stat_mech/statistical_models/network_models_test.py
from absl.testing import absltest
from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import network_models
from epi_forecast_stat_mech.statistical_models import tree_util
from jax.config import config
import jax
import jax.numpy as jnp

    """Tests that `log_likelihood` and `predict` return expected shapes."""
    rng = jax.random.PRNGKey(42)
    c_rng, o_rng, init_rng = jax.random.split(rng, 3)
    covariates = jax.random.uniform(c_rng, (num_locations, num_covariates))
    observations = jax.random.uniform(o_rng, (num_locations, num_observations))
    if dict_keys is not None:
      observations = {key: observations for key in dict_keys}

    if error_model == 'full':
      scale_eps = 1E-2
    else:
      # The length of this needs to agree with the flattened number of
      # "observations" per row. I.e. flattened observations is (7, 6)
      # for the 'plugin' test because the two dict keys double the
      # three observations per location.
      scale_eps = jnp.asarray([.1, .2, .3, .4, .5, .6])
    stat_model = network_models.NormalDistributionModel(
        predict_module=predict_module,
        error_model=error_model,
        scale_eps=scale_eps)
    params = stat_model.init_parameters(init_rng, covariates, observations)

    log_likelihood = stat_model.log_likelihood(params, covariates, observations)

    if error_model == 'full':
      actual_shapes = tree_util.tree_map(lambda x: x.shape, log_likelihood)
      expected_shapes = tree_util.tree_map(lambda x: x.shape, observations)
      self.assertEqual(actual_shapes, expected_shapes)
      distributions_over_observations = stat_model.predict(
          params, covariates, observations)

      actual_shapes = tree_util.tree_map(