code: def create_test_hmm():
      srng = pt.random.RandomStream()
  
      N_pt = pt.iscalar("N")
      N_pt.tag.test_value = 10
      M_pt = pt.iscalar("M")
      M_pt.tag.test_value = 2
  
      mus_pt = pt.matrix("mus")
      mus_pt.tag.test_value = np.stack(
          [np.arange(0.0, 10), np.arange(0.0, -10, -1)], axis=-1
      ).astype(pytensor.config.floatX)
  
      sigmas_pt = pt.ones((N_pt,))
      sigmas_pt.name = "sigmas"
  
      pi_0_rv = srng.dirichlet(pt.ones((M_pt,)), name="pi_0")
      Gamma_rv = srng.dirichlet(pt.ones((M_pt, M_pt)), name="Gamma")
  
      S_0_rv = srng.categorical(pi_0_rv, name="S_0")
  
      def scan_fn(mus_t, sigma_t, S_tm1, Gamma_t):
          S_t = srng.categorical(Gamma_t[S_tm1], name="S_t")
          Y_t = srng.normal(mus_t[S_t], sigma_t, name="Y_t")
          return S_t, Y_t
  
      (S_rv, Y_rv), scan_updates = pytensor.scan(
          fn=scan_fn,
          sequences=[mus_pt, sigmas_pt],
          non_sequences=[Gamma_rv],
          outputs_info=[{"initial": S_0_rv, "taps": [-1]}, {}],
          strict=True,
          name="scan_rv",
      )
      Y_rv.name = "Y_rv"
  
      scan_op = Y_rv.owner.op
      scan_args = ScanArgs.from_node(Y_rv.owner)
  
      Gamma_in = scan_args.inner_in_non_seqs[0]
      Y_t = scan_args.inner_out_nit_sot[0]
      mus_t = scan_args.inner_in_seqs[0]
      sigmas_t = scan_args.inner_in_seqs[1]
      S_t = scan_args.inner_out_sit_sot[0]
      rng_in = scan_args.inner_out_shared[0]
  
      mus_in = Y_rv.owner.inputs[1]
      mus_in.name = "mus_in"
      sigmas_in = Y_rv.owner.inputs[2]
      sigmas_in.name = "sigmas_in"
  
      # The output `S_rv` is really `S_rv[1:]`, so we have to extract the actual
      # `Scan` output: `S_rv`.
      S_in = S_rv.owner.inputs[0]
      S_in.name = "S_in"
  
      return locals()
  
file_path: pytensor/tests/scan/test_utils.py
imports: from copy import copy
  import numpy
  import pytest
  import pytensor
  from pytensor import tensor
  from pytensor.scan.utils import ScanArgs
label_desc: Use image node_id 2 for calling a global function with example usage: create_test_hmm() and returns: locals, S_t,  Y_t