# 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
# Lint as: python3
"""Tests for epi_forecast_stat_mech.statistical_models.network_models.py."""


config.parse_flags_with_absl()  # Necessary for running on TPU.


class NetworkModelsTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(predict_module=network_models.PerceptronModule,
           dict_keys=None,
           num_locations=7,
           num_covariates=5,
           num_observations=3,
           error_model='full'),

# 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

# Lint as: python3
"""Tests for epi_forecast_stat_mech.statistical_models.network_models.py."""


config.parse_flags_with_absl()  # Necessary for running on TPU.


class NetworkModelsTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(predict_module=network_models.PerceptronModule,
           dict_keys=None,
           num_locations=7,
           num_covariates=5,
           num_observations=3,
           error_model='full'),

      dict(predict_module=network_models.LinearModule,