code: def vectorized_model(args, data):
      # Sample global parameters.
      rate_s, prob_i, rho = global_model(args.population)
  
      # Sample reparameterizing variables.
      S_aux = pyro.sample(
          "S_aux",
          dist.Uniform(-0.5, args.population + 0.5)
          .mask(False)
          .expand(data.shape)
          .to_event(1),
      )
      I_aux = pyro.sample(
          "I_aux",
          dist.Uniform(-0.5, args.population + 0.5)
          .mask(False)
          .expand(data.shape)
          .to_event(1),
      )
  
      # Manually enumerate.
      S_curr, S_logp = quantize_enumerate(
          S_aux, min=0, max=args.population
      )
      I_curr, I_logp = quantize_enumerate(
          I_aux, min=0, max=args.population
      )
      # Truncate final value from the right then pad initial value onto the left.
      S_prev = torch.nn.functional.pad(
          S_curr[:-1], (0, 0, 1, 0), value=args.population - 1
      )
      I_prev = torch.nn.functional.pad(
          I_curr[:-1], (0, 0, 1, 0), value=1
      )
      # Reshape to support broadcasting, similar to EnumMessenger.
      T = len(data)
      Q = 4
      S_prev = S_prev.reshape(T, Q, 1, 1, 1)
      I_prev = I_prev.reshape(T, 1, Q, 1, 1)
      S_curr = S_curr.reshape(T, 1, 1, Q, 1)
      S_logp = S_logp.reshape(T, 1, 1, Q, 1)
      I_curr = I_curr.reshape(T, 1, 1, 1, Q)
      I_logp = I_logp.reshape(T, 1, 1, 1, Q)
      data = data.reshape(T, 1, 1, 1, 1)
  
      # Reverse the S2I,I2R computation.
      S2I = S_prev - S_curr
      I2R = I_prev - I_curr + S2I
  
      # Compute probability factors.
      S2I_logp = dist.ExtendedBinomial(
          S_prev, -(rate_s * I_prev).expm1()
      ).log_prob(S2I)
      I2R_logp = dist.ExtendedBinomial(I_prev, prob_i).log_prob(I2R)
      obs_logp = dist.ExtendedBinomial(S2I, rho).log_prob(data)
  
      # Manually perform variable elimination.
      logp = S_logp + (I_logp + obs_logp) + S2I_logp + I2R_logp
      logp = logp.reshape(-1, Q * Q, Q * Q)
      logp = pyro.distributions.hmm._sequential_logmatmulexp(logp)
      logp = logp.reshape(-1).logsumexp(0)
      logp = logp - math.log(
          4
      )  # Account for S,I initial distributions.
      warn_if_nan(logp)
      pyro.factor("obs", logp)
  
imports: import argparse
  import logging
  import math
  import re
  from collections import OrderedDict
  import torch
  import pyro
  import pyro.distributions
  import pyro.distributions.hmm
  import pyro.poutine
  from pyro.infer import MCMC, NUTS, config_enumerate, infer_discrete
  from pyro.infer.autoguide import init_to_value
  from pyro.ops.special import safe_log
  from pyro.ops.tensor_utils import convolve
  from pyro.util import warn_if_nan
label_desc: Use image node_id 12 for calling a global function with example usage: vectorized_model(args, data) without return types