# Path: epi_forecast_stat_mech/evaluation/metrics.py
# def quantile(trajectories, sample_axis=1, quantiles=np.linspace(0, 1, 11)):
# def target_quantile(trajectories, target, sample_axis=1):


import functools
import itertools
import jax
import numpy as np
import tensorflow_probability
from absl.testing import parameterized
from absl.testing import absltest
from epi_forecast_stat_mech.evaluation import metrics
from jax.config import config
          rollout_shape=(101,),
          expected_shape=(11, 101),
          seed=1),
  )
  def testMetricsShape(
      self, metric, batch_size, nsamples, rollout_shape, expected_shape, seed):
    trajectories = jax.random.normal(jax.random.PRNGKey(seed),
                                     (batch_size, nsamples) + rollout_shape)
    metric_value = metric(trajectories)
    self.assertEqual(expected_shape, metric_value.shape)

  @parameterized.parameters(
      dict(batch_size=23,
           nsamples=1000,
           rollout_shape=(11,),
           quantiles=np.linspace(0, 1, 11),
           seed=0),
      dict(batch_size=11,
           nsamples=10000,
           rollout_shape=(3,),
           quantiles=np.linspace(0, 1, 7),
           seed=1)
  )
  def testQuantileAndTargetQuantile(
      self, batch_size, nsamples, rollout_shape, quantiles, seed):
    trajectories = jax.random.normal(jax.random.PRNGKey(seed),
                                     (batch_size, nsamples) + rollout_shape)
    targets = metrics.quantile(trajectories, quantiles=quantiles)
    for target, quantile in zip(targets, quantiles):
      np.testing.assert_allclose(quantile,

# Path: epi_forecast_stat_mech/evaluation/metrics.py
# def quantile(trajectories, sample_axis=1, quantiles=np.linspace(0, 1, 11)):
# def target_quantile(trajectories, target, sample_axis=1):


# Path: epi_forecast_stat_mech/evaluation/metrics_test.py
import functools
import itertools
import jax
import numpy as np
import tensorflow_probability
from absl.testing import parameterized
from absl.testing import absltest
from epi_forecast_stat_mech.evaluation import metrics
from jax.config import config

          rollout_shape=(101,),
          expected_shape=(11, 101),
          seed=1),
  )
  def testMetricsShape(
      self, metric, batch_size, nsamples, rollout_shape, expected_shape, seed):
    trajectories = jax.random.normal(jax.random.PRNGKey(seed),
                                     (batch_size, nsamples) + rollout_shape)
    metric_value = metric(trajectories)
    self.assertEqual(expected_shape, metric_value.shape)

  @parameterized.parameters(
      dict(batch_size=23,
           nsamples=1000,
           rollout_shape=(11,),
           quantiles=np.linspace(0, 1, 11),
           seed=0),
      dict(batch_size=11,
           nsamples=10000,
           rollout_shape=(3,),
           quantiles=np.linspace(0, 1, 7),
           seed=1)
  )
  def testQuantileAndTargetQuantile(
      self, batch_size, nsamples, rollout_shape, quantiles, seed):
    trajectories = jax.random.normal(jax.random.PRNGKey(seed),
                                     (batch_size, nsamples) + rollout_shape)
    targets = metrics.quantile(trajectories, quantiles=quantiles)
    for target, quantile in zip(targets, quantiles):
      np.testing.assert_allclose(quantile,

                                 metrics.target_quantile(trajectories, target),