code: def test_fgraph_rewrite(non_centered_rewrite):
      """Test we can apply a simple rewrite to a PyMC Model."""
  
      with pm.Model(coords={"subject": range(10)}) as m_old:
          group_mean = pm.Normal("group_mean")
          group_std = pm.HalfNormal("group_std")
          subject_mean = pm.Normal(
              "subject_mean", group_mean, group_std, dims=("subject",)
          )
          obs = pm.Normal(
              "obs",
              subject_mean,
              1,
              observed=np.zeros(10),
              dims=("subject",),
          )
  
      fg, _ = fgraph_from_model(m_old)
      non_centered_rewrite.apply(fg)
  
      m_new = model_from_fgraph(fg)
      assert m_new.named_vars_to_dims == {
          "subject_mean": ["subject"],
          "subject_mean_raw_": ["subject"],
          "obs": ["subject"],
      }
      assert set(m_new.named_vars) == {
          "group_mean",
          "group_std",
          "subject_mean_raw_",
          "subject_mean",
          "obs",
      }
      assert {rv.name for rv in m_new.free_RVs} == {
          "group_mean",
          "group_std",
          "subject_mean_raw_",
      }
      assert {rv.name for rv in m_new.observed_RVs} == {"obs"}
      assert {rv.name for rv in m_new.deterministics} == {
          "subject_mean"
      }
  
      with pm.Model() as m_ref:
          group_mean = pm.Normal("group_mean")
          group_std = pm.HalfNormal("group_std")
          subject_mean_raw = pm.Normal(
              "subject_mean_raw_", 0, 1, shape=(10,)
          )
          subject_mean = pm.Deterministic(
              "subject_mean", group_mean + subject_mean_raw * group_std
          )
          obs = pm.Normal("obs", subject_mean, 1, observed=np.zeros(10))
  
      np.testing.assert_array_equal(
          pm.draw(m_new["subject_mean_raw_"], draws=7, random_seed=1),
          pm.draw(m_ref["subject_mean_raw_"], draws=7, random_seed=1),
      )
  
      ip = m_new.initial_point()
      np.testing.assert_equal(
          m_new.compile_logp()(ip),
          m_ref.compile_logp()(ip),
      )
  
file_path: pymc/tests/model/test_fgraph.py
imports: import numpy
  import pytensor.tensor
  import pytest
  from pytensor import config, shared
  from pytensor.graph import Constant, FunctionGraph, node_rewriter
  from pytensor.graph.rewriting.basic import in2out
  from pytensor.tensor.exceptions import NotScalarConstantError
  import pymc
  from pymc.model.fgraph import ModelDeterministic, ModelFreeRV, ModelNamed, ModelObservedRV, ModelPotential, ModelVar, clone_model, fgraph_from_model, model_deterministic, model_free_rv, model_from_fgraph
label_desc: Use image node_id 9 for calling a global function with example usage: test_fgraph_rewrite(non_centered_rewrite) without return types