code: def test_hetero_mixture_categorical(
      X_args,
      Y_args,
      Z_args,
      p_val,
      comp_size,
      idx_size,
      extra_indices,
      join_axis,
      supported,
  ):
      X_rv = pt.random.normal(*X_args, size=comp_size, name="X")
      Y_rv = pt.random.gamma(
          Y_args[0], scale=Y_args[1], size=comp_size, name="Y"
      )
      Z_rv = pt.random.normal(*Z_args, size=comp_size, name="Z")
  
      p_at = pt.as_tensor(p_val).type()
      p_at.name = "p"
      p_at.tag.test_value = np.array(
          p_val, dtype=pytensor.config.floatX
      )
      I_rv = pt.random.categorical(p_at, size=idx_size, name="I")
  
      i_vv = I_rv.clone()
      i_vv.name = "i"
  
      indices_at = list(extra_indices)
      indices_at.insert(join_axis, I_rv)
      indices_at = tuple(indices_at)
  
      M_rv = pt.stack([X_rv, Y_rv, Z_rv], axis=join_axis)[indices_at]
      M_rv.name = "M"
  
      m_vv = M_rv.clone()
      m_vv.name = "m"
  
      if supported:
          logp_parts = conditional_logp(
              {M_rv: m_vv, I_rv: i_vv}, sum=False
          )
      else:
          with pytest.raises(
              RuntimeError, match="could not be derived: {m}"
          ):
              conditional_logp({M_rv: m_vv, I_rv: i_vv}, sum=False)
          return
  
      I_logp_fn = pytensor.function([p_at, i_vv], logp_parts[i_vv])
      M_logp_fn = pytensor.function([m_vv, i_vv], logp_parts[m_vv])
  
      assert_no_rvs(I_logp_fn.maker.fgraph.outputs[0])
      assert_no_rvs(M_logp_fn.maker.fgraph.outputs[0])
  
      decimals = 6 if pytensor.config.floatX == "float64" else 4
  
      test_val_rng = np.random.RandomState(3238)
  
      norm_1_sp = sp.norm(loc=X_args[0], scale=X_args[1])
      gamma_sp = sp.gamma(Y_args[0], scale=Y_args[1])
      norm_2_sp = sp.norm(loc=Z_args[0], scale=Z_args[1])
  
      # Handle scipy annoying squeeze of random draws
      real_comp_size = tuple(X_rv.shape.eval())
  
      for i in range(10):
          i_val = CategoricalRV.rng_fn(test_val_rng, p_val, idx_size)
  
          indices_val = list(extra_indices)
          indices_val.insert(join_axis, i_val)
          indices_val = tuple(indices_val)
  
          x_val = np.broadcast_to(
              norm_1_sp.rvs(size=comp_size, random_state=test_val_rng),
              real_comp_size,
          )
          y_val = np.broadcast_to(
              gamma_sp.rvs(size=comp_size, random_state=test_val_rng),
              real_comp_size,
          )
          z_val = np.broadcast_to(
              norm_2_sp.rvs(size=comp_size, random_state=test_val_rng),
              real_comp_size,
          )
  
          component_logps = np.stack(
              [
                  norm_1_sp.logpdf(x_val),
                  gamma_sp.logpdf(y_val),
                  norm_2_sp.logpdf(z_val),
              ],
              axis=join_axis,
          )[indices_val]
          index_logps = scipy_logprob(i_val, p_val)
          exp_obs_logps = (
              component_logps
              + index_logps[(Ellipsis,) + (None,) * join_axis]
          )
  
          m_val = np.stack([x_val, y_val, z_val], axis=join_axis)[
              indices_val
          ]
  
          I_logp_vals = I_logp_fn(p_val, i_val)
          M_logp_vals = M_logp_fn(m_val, i_val)
  
          logp_vals = (
              M_logp_vals
              + I_logp_vals[(Ellipsis,) + (None,) * join_axis]
          )
  
          np.testing.assert_almost_equal(
              logp_vals, exp_obs_logps, decimal=decimals
          )
  
imports: import numpy
  import pytensor
  import pytensor.tensor
  import pytest
  import scipy.stats.distributions
  from pytensor import function
  from pytensor.graph.basic import Variable
  from pytensor.ifelse import ifelse
  from pytensor.tensor.random.basic import CategoricalRV
  from pytensor.tensor.shape import shape_tuple
  from pytensor.tensor.subtensor import AdvancedSubtensor, AdvancedSubtensor1, Subtensor, as_index_constant
  from pymc.logprob.abstract import MeasurableVariable
  from pymc.logprob.basic import conditional_logp, logp
  from pymc.logprob.mixture import MeasurableSwitchMixture, expand_indices
  from pymc.logprob.rewriting import construct_ir_fgraph
  from pymc.logprob.utils import dirac_delta
  from pymc.testing import assert_no_rvs
  from tests.logprob.utils import scipy_logprob
label_desc: Use image node_id 4 for calling a global function with example usage: test_hetero_mixture_categorical(X_args, Y_args, Z_args, p_val, comp_size, idx_size, extra_indices, join_axis, supported) without return types