code: def pyrocov_model(dataset):
      # Tensor shapes are commented at the end of some lines.
      features = dataset["features"]
      local_time = dataset["local_time"][..., None]  # [T, P, 1]
      T, P, _ = local_time.shape
      S, F = features.shape
      weekly_strains = dataset["weekly_strains"]
      assert weekly_strains.shape == (T, P, S)
  
      # Sample global random variables.
      coef_scale = pyro.sample(
          "coef_scale", dist.InverseGamma(5e3, 1e2)
      )[..., None]
      rate_scale = pyro.sample("rate_scale", dist.LogNormal(-4, 2))[
          ..., None
      ]
      init_loc_scale = pyro.sample(
          "init_loc_scale", dist.LogNormal(0, 2)
      )[..., None]
      init_scale = pyro.sample("init_scale", dist.LogNormal(0, 2))[
          ..., None
      ]
  
      # Assume relative growth rate depends strongly on mutations and weakly on place.
      coef_loc = torch.zeros(F)
      coef = pyro.sample(
          "coef", dist.Logistic(coef_loc, coef_scale).to_event(1)
      )  # [F]
      rate_loc = pyro.deterministic(
          "rate_loc", 0.01 * coef @ features.T, event_dim=1
      )  # [S]
  
      # Assume initial infections depend strongly on strain and place.
      init_loc = pyro.sample(
          "init_loc",
          dist.Normal(torch.zeros(S), init_loc_scale).to_event(1),
      )  # [S]
      with pyro.plate("place", P, dim=-1):
          rate = pyro.sample(
              "rate", dist.Normal(rate_loc, rate_scale).to_event(1)
          )  # [P, S]
          init = pyro.sample(
              "init", dist.Normal(init_loc, init_scale).to_event(1)
          )  # [P, S]
  
          # Finally observe counts.
          with pyro.plate("time", T, dim=-2):
              logits = init + rate * local_time  # [T, P, S]
              pyro.sample(
                  "obs",
                  dist.Multinomial(logits=logits, validate_args=False),
                  obs=weekly_strains,
              )
  
imports: from collections import OrderedDict, namedtuple
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer import SVI, JitTrace_ELBO, Predictive, Trace_ELBO
  from pyro.infer.autoguide import AutoGaussian, AutoGuideList
  from pyro.infer.autoguide.gaussian import AutoGaussianDense, AutoGaussianFunsor, _break_plates
  from pyro.infer.reparam import LocScaleReparam
  from pyro.optim import ClippedAdam
  from tests.common import assert_close, assert_equal, xfail_if_not_implemented
label_desc: Use image node_id 16 for calling a global function with example usage: pyrocov_model(dataset) without return types