explore using pymc for the defined distributions instead

[kratsg]

import pymc as pm
import pytensor.tensor as pt

# Define a random variable
with pm.Model():
    mu = pm.Normal('mu', mu=0, sigma=1) 
    x = pm.Normal('x', mu=mu, sigma=1)

# Create a symbolic value for x
x_value = pt.scalar('x_value')

# Get the log-probability graph
log_prob_graph = pm.logp(x, x_value)

# You can evaluate the log-probability for a specific value using .eval()
# (useful for debugging, but not for inference)
print(log_prob_graph.eval({x_value: 0.5, mu: 0.0})) 

# For repeated evaluations, compile the function
log_prob_fn = pm.compile_pymc([x_value, mu], log_prob_graph)
print(log_prob_fn(x_value=0.5, mu=0.0))

with something like

>>> import pymc
>>> pymc.distributions.Normal
<class 'pymc.distributions.continuous.Normal'>
>>> import pytensor.tensor as pt
>>> with pymc.Model():
...     foo = pymc.distributions.Normal("test", mu=pt.scalar("mu"), sigma=pt.scalar("sigma"))
...
>>> foo
test
>>> foo.name
'test'
>>> type(foo)
<class 'pytensor.tensor.variable.TensorVariable'>
>>> import pytensor
>>> pytensor.pp(foo)
'normal_rv{"(),()->()"}(RNG(<Generator(PCG64) at 0x134EC82E0>), None, mu, sigma)'

>>> pymc.logp(foo, 1.0)
test_logprob
>>> pytensor.pp(pymc.logp(foo, 1.0))
'Check{sigma > 0}((((-0.5 * (((1.0 - mu) / sigma) ** 2)) - Log(Sqrt(6.283185307179586))) - Log(sigma)), All{axes=None}([All{axes=None}(Gt(sigma, 0))]))'
>>> pytensor.pp(pymc.logp(foo, [1.0, 2.0]))
'Check{sigma > 0}((((-0.5 * ((([1. 2.] - mu) / sigma) ** 2)) - Log(Sqrt(6.283185307179586))) - Log(sigma)), All{axes=None}([All{axes=None}(Gt(sigma, 0))]))'
>>> pytensor.pp(pymc.logp(foo, [1.0, 2.0]))

Useful links:

• https://github.com/pymc-devs/pytensor/blob/main/pytensor/tensor/random/basic.py
• https://github.com/pymc-devs/pymc/blob/main/pymc/distributions/distribution.py#L417
• https://github.com/pymc-devs/pymc/blob/main/pymc/distributions/continuous.py
• https://github.com/pymc-devs/pymc/blob/main/pymc/distributions/distribution.py#L417