BUG: Unexpected TypeError when inverting scipy-based function

[Spiruel]

Describe the issue:

I am adding parameters into a model in an attempt to perform bayesian inversion. I am getting a typeerror when introducing variables, which is unexpected behaviour.

Reproduceable code example:

import numpy as np
import pymc as pm
import pytensor.tensor as pt
from pytensor.graph import Op, Apply
import arviz as az
import prosail

# Define the new model function with additional parameters
def my_model(n, cab, car, cbrown, cw, cm, lai, tto, tts, x):
    reflectance = prosail.run_prosail(n=n, cab=cab, car=car, cbrown=cbrown, cw=cw, cm=cm, lai=lai, tto=10., tts=30., lidfa=-1.0, hspot=0.01, psi=0., prospect_version="D", factor='SDR', rsoil=.5, psoil=.5)
    return reflectance

def my_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data):
    for param in (n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data):
        if not isinstance(param, (float, np.ndarray)):
            raise TypeError(f"Invalid input type to loglike: {type(param)}")
    model = my_model(n, cab, car, cbrown, cw, cm, lai, tto, tts, x)
    return -0.5 * ((data - model) / sigma) ** 2 - np.log(np.sqrt(2 * np.pi)) - np.log(sigma)

class LogLike(Op):
    def make_node(self, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data) -> Apply:
        n = pt.as_tensor(n)
        cab = pt.as_tensor(cab)
        car = pt.as_tensor(car)
        cbrown = pt.as_tensor(cbrown)
        cw = pt.as_tensor(cw)
        cm = pt.as_tensor(cm)
        lai = pt.as_tensor(lai)
        tto = pt.as_tensor(tto)
        tts = pt.as_tensor(tts)
        sigma = pt.as_tensor(sigma)
        x = pt.as_tensor(x)
        data = pt.as_tensor(data)

        inputs = [n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data]
        outputs = [data.type()]

        return Apply(self, inputs, outputs)

    def perform(self, node: Apply, inputs: list[np.ndarray], outputs: list[list[None]]) -> None:
        n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data = inputs
        loglike_eval = my_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data)
        outputs[0][0] = np.asarray(loglike_eval)

def custom_dist_loglike(data, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x):
    return loglike_op(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data)

N = 2101
sigma = 1.0
x = np.linspace(0.0, 9.0, N)

# True parameter values for testing
ntrue = 0.4
cabtrue = 0.1
cartrue = 0.2
cbrowntrue = 0.3
cwtrue = 0.05
cmtrue = 0.07
laitrue = 3.0
ttotrue = 0.15
ttstrue = 0.25

truemodel = my_model(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, laitrue, ttotrue, ttstrue, x)

rng = np.random.default_rng(716743)
data = sigma * rng.normal(size=N) + truemodel

loglike_op = LogLike()

test_out = loglike_op(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, laitrue, ttotrue, ttstrue, sigma, x, data)

with pm.Model() as no_grad_model:

    n = pm.Uniform("n", lower=-1.0, upper=5.0, initval=ntrue)
    cab = pm.Uniform("cab", lower=-5, upper=60, initval=cabtrue)
    car = pm.Uniform("car", lower=-5, upper=20, initval=cartrue)
    cbrown = pm.Uniform("cbrown", lower=0, upper=2, initval=cbrowntrue)
    cw = pm.Uniform("cw", lower=-10.0, upper=10, initval=cwtrue)
    cm = pm.Uniform("cm", lower=0, upper=0.2, initval=cmtrue)
    lai = pm.Uniform("lai", lower=0.5, upper=5.0, initval=laitrue)
    tto = pm.Uniform("tto", lower=0.0, upper=20.0, initval=ttotrue)
    tts = pm.Uniform("tts", lower=0, upper=60.0, initval=ttstrue)

    likelihood = pm.CustomDist(
        "likelihood", n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, observed=data, logp=custom_dist_loglike
    )

import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore")

    with no_grad_model:
        idata_no_grad = pm.sample(100000, step=pm.DEMetropolisZ())

Error message:

TypeError: No matching definition for argument type(s) float64, array(float64, 0d, C), float64, float64
Apply node that caused the error: LogLike(Composite{((i2 * sigmoid(i0)) - (i1 - sigmoid(i0)))}.0, Composite{((i2 * sigmoid(i0)) + (i3 * (i1 - sigmoid(i0))))}.0, Composite{((i2 * sigmoid(i0)) + (i3 * (i1 - sigmoid(i0))))}.0, Composite{(i1 * sigmoid(i0))}.0, Composite{((i2 * sigmoid(i0)) + (i3 * (i1 - sigmoid(i0))))}.0, Composite{(i1 * sigmoid(i0))}.0, Composite{((i2 * sigmoid(i0)) + (i3 * (i1 - sigmoid(i0))))}.0, Composite{(i1 * sigmoid(i0))}.0, Composite{(i1 * sigmoid(i0))}.0, 1.0, [0.0000000 ... 00000e+00], likelihood{[-0.589307 ... .02688868]})
Toposort index: 9
Inputs types: [TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float64, shape=()), TensorType(float32, shape=()), TensorType(float64, shape=(2101,)), TensorType(float64, shape=(2101,))]
Inputs shapes: [(), (), (), (), (), (), (), (), (), (), (2101,), (2101,)]
Inputs strides: [(), (), (), (), (), (), (), (), (), (), (8,), (8,)]
Inputs values: [array(0.4), array(0.1), array(0.2), array(0.3), array(0.05), array(0.07), array(3.), array(0.15), array(0.25), array(1., dtype=float32), 'not shown', 'not shown']
Inputs type_num: [12, 12, 12, 12, 12, 12, 12, 12, 12, 11, 12, 12]
Outputs clients: [['output']]

PyMC version information:

'5.12.0'

Context for the issue:

No response

[welcome]

]
🎉 Welcome to PyMC! 🎉 We're really excited to have your input into the project! 💖

If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.

[ricardoV94]

The error seems to happen inside your Loglike code? Doesn't look like a PyMC or PyTensor error