code: def test_scan_joint_logprob(require_inner_rewrites):
      srng = pt.random.RandomStream()
  
      N_tt = pt.iscalar("N")
      N_val = 10
      N_tt.tag.test_value = N_val
  
      M_tt = pt.iscalar("M")
      M_val = 2
      M_tt.tag.test_value = M_val
  
      mus_tt = pt.matrix("mus_t")
  
      mus_val = np.stack(
          [np.arange(0.0, 10), np.arange(0.0, -10, -1)], axis=-1
      ).astype(pytensor.config.floatX)
      mus_tt.tag.test_value = mus_val
  
      sigmas_tt = pt.ones((N_tt,))
      Gamma_rv = srng.dirichlet(pt.ones((M_tt, M_tt)), name="Gamma")
  
      Gamma_vv = Gamma_rv.clone()
      Gamma_vv.name = "Gamma_vv"
  
      Gamma_val = np.array([[0.5, 0.5], [0.5, 0.5]])
      Gamma_rv.tag.test_value = Gamma_val
  
      def scan_fn(mus_t, sigma_t, Gamma_t):
          S_t = srng.categorical(Gamma_t[0], name="S_t")
  
          if require_inner_rewrites:
              Y_t = srng.normal(mus_t, sigma_t, name="Y_t")[S_t]
          else:
              Y_t = srng.normal(mus_t[S_t], sigma_t, name="Y_t")
  
          return Y_t, S_t
  
      (Y_rv, S_rv), _ = pytensor.scan(
          fn=scan_fn,
          sequences=[mus_tt, sigmas_tt],
          non_sequences=[Gamma_rv],
          outputs_info=[{}, {}],
          strict=True,
          name="scan_rv",
      )
      Y_rv.name = "Y"
      S_rv.name = "S"
  
      y_vv = Y_rv.clone()
      y_vv.name = "y"
  
      s_vv = S_rv.clone()
      s_vv.name = "s"
  
      y_logp = conditional_logp(
          {Y_rv: y_vv, S_rv: s_vv, Gamma_rv: Gamma_vv}
      )
      y_logp_combined = pt.sum(
          [pt.sum(factor) for factor in y_logp.values()]
      )
  
      y_val = np.arange(10)
      s_val = np.array([0, 1, 0, 1, 1, 0, 0, 0, 1, 1])
  
      test_point = {
          y_vv: y_val,
          s_vv: s_val,
          M_tt: M_val,
          N_tt: N_val,
          mus_tt: mus_val,
          Gamma_vv: Gamma_val,
      }
  
      y_logp_fn = pytensor.function(
          list(test_point.keys()), y_logp_combined
      )
  
      assert_no_rvs(y_logp_fn.maker.fgraph.outputs[0])
  
      # Construct the joint log-probability by hand so we can compare it with
      # `y_logp`
      def scan_fn(mus_t, sigma_t, Y_t_val, S_t_val, Gamma_t):
          S_t = pt.random.categorical(Gamma_t[0], name="S_t")
          Y_t = pt.random.normal(mus_t[S_t_val], sigma_t, name="Y_t")
          Y_t_logp, S_t_logp = _logprob_helper(
              Y_t, Y_t_val
          ), _logprob_helper(S_t, S_t_val)
          Y_t_logp.name = "log(Y_t=y_t)"
          S_t_logp.name = "log(S_t=s_t)"
          return Y_t_logp, S_t_logp
  
      (Y_rv_logp, S_rv_logp), _ = pytensor.scan(
          fn=scan_fn,
          sequences=[mus_tt, sigmas_tt, y_vv, s_vv],
          non_sequences=[Gamma_vv],
          outputs_info=[{}, {}],
          strict=True,
          name="scan_rv",
      )
      Y_rv_logp.name = "logp(Y=y)"
      S_rv_logp.name = "logp(S=s)"
  
      Gamma_logp = _logprob_helper(Gamma_rv, Gamma_vv)
  
      y_logp_ref = Y_rv_logp.sum() + S_rv_logp.sum() + Gamma_logp.sum()
  
      assert_no_rvs(y_logp_ref)
  
      y_logp_val = y_logp_combined.eval(test_point)
  
      y_logp_ref_val = y_logp_ref.eval(test_point)
  
      assert np.allclose(y_logp_val, y_logp_ref_val)
  
imports: import numpy
  import pytensor
  import pytensor.tensor
  import pytest
  from pytensor import Mode
  from pytensor.raise_op import assert_op
  from pytensor.scan.utils import ScanArgs
  from scipy import stats
  from pymc.logprob.abstract import _logprob_helper
  from pymc.logprob.basic import conditional_logp, logp
  from pymc.logprob.scan import construct_scan, convert_outer_out_to_in, get_random_outer_outputs
  from pymc.testing import assert_no_rvs
label_desc: Use image node_id 4 for calling a global function with example usage: test_scan_joint_logprob(require_inner_rewrites) and returns: Y_t,  S_t, Y_t_logp,  S_t_logp