numba: Incorrect gradient and typing error

[aseyboldt]

Description

The following shows two different results with the C-backend and the numba backend:

import pymc as pm
import pytensor.tensor as pt
import numpy as np
import pytensor

np.random.seed(1)

n = 100
X = np.linspace(0, 10, n)[:, None]
y = np.random.randn(n)

with pm.Model(coords={"x_dim": X[:, 0], "x_dim2": [0]}) as model:
    ell = pm.InverseGamma("ell", mu=1.0, sigma=0.6)
    eta = pm.Exponential("eta", scale=3.0)

    cov = eta**2 * pm.gp.cov.Matern52(1, ls=ell)
    gp = pm.gp.Marginal(cov_func=cov)

    sigma = pm.HalfCauchy("sigma", beta=5)

    # Typing error without this!!!
    X = pm.Data("X_data", X)
    y = pm.Data("y_data", y)
    
    y_ = gp.marginal_likelihood("y", X=X, y=y, sigma=sigma)


logp = model.logp()
grad = pt.grad(logp, model.value_vars)
func1 = pytensor.function(model.value_vars, [logp, *grad], mode="NUMBA")
func2 = pytensor.function(model.value_vars, [logp, *grad], mode="FAST_RUN")
point = model.initial_point()
print(func1(**point))
print(func2(**point))

[array(-263.86874314), array(-16.06796754), array(-121.94404258), array(24.92776397)]
[array(-263.86874314), array(1.97580188), array(-7.71406399), array(-89.30221005)]

In addition, without the pm.Data calls I get typeing errors:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of type(CPUDispatcher(<function careduce_axis at 0x7eee5b6f72e0>)) with parameters (readonly array(float64, 2d, C))
Known signatures:
 * (Array(float64, 2, 'A', False, aligned=True),) -> array(float64, 1d, A)
During: resolving callee type: type(CPUDispatcher(<function careduce_axis at 0x7eee5b6f72e0>))
During: typing of call at [/tmp/tmpfu_6ldpg](http://localhost:7890/tmp/tmpfu_6ldpg) (89)


File "../../../../tmp/tmpfu_6ldpg", line 89:
def numba_funcified_fgraph(ell_log_, eta_log_, sigma_log_):
    <source elided>
    # Sum{axis=1}(Sqr.0)
    tensor_variable_55 = careduce_axis_1(tensor_variable_53)