code: def pyrocov_model_poisson(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"]  # [T, P, S]
      if not torch._C._get_tracing_state():
          assert weekly_strains.shape == (T, P, S)
      strain_plate = pyro.plate("strain", S, dim=-1)
      place_plate = pyro.plate("place", P, dim=-2)
      time_plate = pyro.plate("time", T, dim=-3)
  
      # Sample global random variables.
      coef_scale = pyro.sample("coef_scale", dist.LogNormal(-4, 2))
      rate_loc_scale = pyro.sample(
          "rate_loc_scale", dist.LogNormal(-4, 2)
      )
      rate_scale = pyro.sample("rate_scale", dist.LogNormal(-4, 2))
      init_loc_scale = pyro.sample(
          "init_loc_scale", dist.LogNormal(0, 2)
      )
      init_scale = pyro.sample("init_scale", dist.LogNormal(0, 2))
      pois_loc = pyro.sample("pois_loc", dist.Normal(0, 2))
      pois_scale = pyro.sample("pois_scale", dist.LogNormal(0, 2))
  
      coef = pyro.sample(
          "coef", dist.Logistic(torch.zeros(F), coef_scale).to_event(1)
      )  # [F]
      rate_loc_loc = 0.01 * coef @ features.T
      with strain_plate:
          rate_loc = pyro.sample(
              "rate_loc", dist.Normal(rate_loc_loc, rate_loc_scale)
          )  # [S]
          init_loc = pyro.sample(
              "init_loc", dist.Normal(0, init_loc_scale)
          )  # [S]
      with place_plate, strain_plate:
          rate = pyro.sample(
              "rate", dist.Normal(rate_loc, rate_scale)
          )  # [P, S]
          init = pyro.sample(
              "init", dist.Normal(init_loc, init_scale)
          )  # [P, S]
  
      # Finally observe counts.
      with time_plate, place_plate:
          pois = pyro.sample(
              "pois", dist.LogNormal(pois_loc, pois_scale)
          )
      with time_plate, place_plate, strain_plate:
          # Note .softmax() breaks conditional independence over strain, but only
          # weakly. We could directly call .exp(), but .softmax is more
          # numerically stable.
          logits = pois * (init + rate * local_time).softmax(
              -1
          )  # [T, P, S]
          pyro.sample("obs", dist.Poisson(logits), 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 19 for calling a global function with example usage: pyrocov_model_poisson(dataset) without return types