# Path: epi_forecast_stat_mech/statistical_models/base.py
# class StatisticalModel:
#   def init_parameters(self):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, observations):
#   def predict(self, parameters, covariates):
# 
# Path: epi_forecast_stat_mech/statistical_models/probability.py
# def softplus(x, floor, softness=1):
# def softplus_scale_estimate(
#     x, axis, min_scale, softness=1):
# def gaussian_with_softplus_scale_estimate(
#     x, axis, min_scale, mean=0, softness=1):
# def pseudo_poisson_log_probs(intensity, observations, axis=0):
# def soft_laplace_log_prob(x, scale=1):
# def laplace_prior(parameters, scale_parameter=1.):
# def log_mixed_laplace(x, a=0., b=0.5, include_constants=False):
# def soft_abs(x, delta):
# def log_soft_mixed_laplace(x, a=0., b=0.5, delta=0.01, include_constants=False):
# def log_soft_mixed_laplace_single(x,
#                                   a=0.,
#                                   b=0.5,
#                                   delta=0.01,
#                                   include_constants=False):
# def apply_to_recursive_dict_by_last_key(last_key_mapper, d, default):
# def _rec_apply_to_recursive_dict_by_last_key(last_key_mapper, d, last_key,
#                                              default):
# def log_soft_mixed_laplace_on_kernels(flax_dict):
# def bounded_derivative_log(x, baseline_constant):
# def gaussian_error_logprob_with_bottom_scale(error, bottom_scale, axis=0):
# LOG_SOFT_LAPLACE_INTEGRAL = 1.185495232349193
# 
# 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):


import dataclasses
import jax
import jax.numpy as jnp
import tensorflow_probability as tfp
from typing import Any, Callable
from epi_forecast_stat_mech.statistical_models import base
from epi_forecast_stat_mech.statistical_models import probability as stat_prob
from epi_forecast_stat_mech.statistical_models import tree_util
from flax.deprecated import nn
    Returns:
      pytree of probability distributions over `observations` given the
      `parameters` of the statistical model and `covariates`.
    """
    output_size, unpack_fn = self._output_size_and_unpack_fn(observations)
    raw_predictions = self.predict_module.call(
        parameters, covariates, output_size)
    if self.error_model == 'full':
      loc, raw_scale = jnp.split(raw_predictions, 2, -1)
      scale = jax.nn.softplus(raw_scale) + self.scale_eps
      return unpack_fn(loc), unpack_fn(scale)
    elif self.error_model == 'plugin':
      loc = raw_predictions
      return unpack_fn(loc), None
    else:
      raise ValueError(f'unexpected error_model: {self.error_model}')

  def predict(self, parameters, covariates, observations):
    """Predicts a distribution over `observations` based on `covariates`.

    Args:
      parameters: parameters of the statistical model.
      covariates: array representing covariates for each location.
      observations: structure of observations to be predicted.

    Returns:
      pytree of probability distributions over `observations` given the
      `parameters` of the statistical model and `covariates`.
    """
    loc, scale = self.get_loc_scale(parameters, covariates, observations)

# Path: epi_forecast_stat_mech/statistical_models/base.py
# class StatisticalModel:
#   def init_parameters(self):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, observations):
#   def predict(self, parameters, covariates):
# 
# Path: epi_forecast_stat_mech/statistical_models/probability.py
# def softplus(x, floor, softness=1):
# def softplus_scale_estimate(
#     x, axis, min_scale, softness=1):
# def gaussian_with_softplus_scale_estimate(
#     x, axis, min_scale, mean=0, softness=1):
# def pseudo_poisson_log_probs(intensity, observations, axis=0):
# def soft_laplace_log_prob(x, scale=1):
# def laplace_prior(parameters, scale_parameter=1.):
# def log_mixed_laplace(x, a=0., b=0.5, include_constants=False):
# def soft_abs(x, delta):
# def log_soft_mixed_laplace(x, a=0., b=0.5, delta=0.01, include_constants=False):
# def log_soft_mixed_laplace_single(x,
#                                   a=0.,
#                                   b=0.5,
#                                   delta=0.01,
#                                   include_constants=False):
# def apply_to_recursive_dict_by_last_key(last_key_mapper, d, default):
# def _rec_apply_to_recursive_dict_by_last_key(last_key_mapper, d, last_key,
#                                              default):
# def log_soft_mixed_laplace_on_kernels(flax_dict):
# def bounded_derivative_log(x, baseline_constant):
# def gaussian_error_logprob_with_bottom_scale(error, bottom_scale, axis=0):
# LOG_SOFT_LAPLACE_INTEGRAL = 1.185495232349193
# 
# 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.py
import dataclasses
import jax
import jax.numpy as jnp
import tensorflow_probability as tfp
from typing import Any, Callable
from epi_forecast_stat_mech.statistical_models import base
from epi_forecast_stat_mech.statistical_models import probability as stat_prob
from epi_forecast_stat_mech.statistical_models import tree_util
from flax.deprecated import nn

    Returns:
      pytree of probability distributions over `observations` given the
      `parameters` of the statistical model and `covariates`.
    """
    output_size, unpack_fn = self._output_size_and_unpack_fn(observations)
    raw_predictions = self.predict_module.call(
        parameters, covariates, output_size)
    if self.error_model == 'full':
      loc, raw_scale = jnp.split(raw_predictions, 2, -1)
      scale = jax.nn.softplus(raw_scale) + self.scale_eps
      return unpack_fn(loc), unpack_fn(scale)
    elif self.error_model == 'plugin':
      loc = raw_predictions
      return unpack_fn(loc), None
    else:
      raise ValueError(f'unexpected error_model: {self.error_model}')

  def predict(self, parameters, covariates, observations):
    """Predicts a distribution over `observations` based on `covariates`.

    Args:
      parameters: parameters of the statistical model.
      covariates: array representing covariates for each location.
      observations: structure of observations to be predicted.

    Returns:
      pytree of probability distributions over `observations` given the
      `parameters` of the statistical model and `covariates`.
    """
    loc, scale = self.get_loc_scale(parameters, covariates, observations)

    return tree_util.tree_multimap(