code: def test_scan_SEIR():
      """Test a scan implementation of a SEIR model.
  
      SEIR model definition:
      S[t+1] = S[t] - B[t]
      E[t+1] = E[t] +B[t] - C[t]
      I[t+1] = I[t+1] + C[t] - D[t]
  
      B[t] ~ Binom(S[t], beta)
      C[t] ~ Binom(E[t], gamma)
      D[t] ~ Binom(I[t], delta)
      """
  
      def binomln(n, k):
          return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)
  
      def binom_log_prob(n, p, value):
          return (
              binomln(n, value)
              + value * log(p)
              + (n - value) * log(1 - p)
          )
  
      # sequences
      at_C = vector("C_t", dtype="int32", shape=(8,))
      at_D = vector("D_t", dtype="int32", shape=(8,))
      # outputs_info (initial conditions)
      st0 = lscalar("s_t0")
      et0 = lscalar("e_t0")
      it0 = lscalar("i_t0")
      logp_c = scalar("logp_c")
      logp_d = scalar("logp_d")
      # non_sequences
      beta = scalar("beta")
      gamma = scalar("gamma")
      delta = scalar("delta")
  
      # TODO: Use random streams when their JAX conversions are implemented.
      # trng = pytensor.tensor.random.RandomStream(1234)
  
      def seir_one_step(
          ct0, dt0, st0, et0, it0, logp_c, logp_d, beta, gamma, delta
      ):
          # bt0 = trng.binomial(n=st0, p=beta)
          bt0 = st0 * beta
          bt0 = bt0.astype(st0.dtype)
  
          logp_c1 = binom_log_prob(et0, gamma, ct0).astype(logp_c.dtype)
          logp_d1 = binom_log_prob(it0, delta, dt0).astype(logp_d.dtype)
  
          st1 = st0 - bt0
          et1 = et0 + bt0 - ct0
          it1 = it0 + ct0 - dt0
          return st1, et1, it1, logp_c1, logp_d1
  
      (st, et, it, logp_c_all, logp_d_all), _ = scan(
          fn=seir_one_step,
          sequences=[at_C, at_D],
          outputs_info=[st0, et0, it0, logp_c, logp_d],
          non_sequences=[beta, gamma, delta],
      )
      st.name = "S_t"
      et.name = "E_t"
      it.name = "I_t"
      logp_c_all.name = "C_t_logp"
      logp_d_all.name = "D_t_logp"
  
      out_fg = FunctionGraph(
          [
              at_C,
              at_D,
              st0,
              et0,
              it0,
              logp_c,
              logp_d,
              beta,
              gamma,
              delta,
          ],
          [st, et, it, logp_c_all, logp_d_all],
      )
  
      s0, e0, i0 = 100, 50, 25
      logp_c0 = np.array(0.0, dtype=config.floatX)
      logp_d0 = np.array(0.0, dtype=config.floatX)
      beta_val, gamma_val, delta_val = (
          np.array(val, dtype=config.floatX)
          for val in [0.277792, 0.135330, 0.108753]
      )
      C = np.array([3, 5, 8, 13, 21, 26, 10, 3], dtype=np.int32)
      D = np.array([1, 2, 3, 7, 9, 11, 5, 1], dtype=np.int32)
  
      test_input_vals = [
          C,
          D,
          s0,
          e0,
          i0,
          logp_c0,
          logp_d0,
          beta_val,
          gamma_val,
          delta_val,
      ]
      compare_jax_and_py(out_fg, test_input_vals)
  
imports: import re
  import numpy
  import pytest
  import pytensor.tensor
  from pytensor import function, shared
  from pytensor.compile import get_mode
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.scan import until
  from pytensor.scan.basic import scan
  from pytensor.scan.op import Scan
  from pytensor.tensor import random
  from pytensor.tensor.math import gammaln, log
  from pytensor.tensor.type import dmatrix, dvector, lscalar, matrix, scalar, vector
  from tests.link.jax.test_basic import compare_jax_and_py
label_desc: Use image node_id 9 for calling a global function with example usage: test_scan_SEIR() and returns: unknown, unknown, st1,  et1,  it1,  logp_c1,  logp_d1