I cannot use any custom kernel to fit the Gp model for fitting in ExactGP

[Saha152]

I have implemented the following custom kernel & prior to use in my inversion, as the existing matern kernel only uses the default parameter values.

import jax.numpy as jnp
from jax import vmap
import numpyro
import numpyro.distributions as dist
#from typing import Dict
from typing import Dict, Callable # Import Callable from typing

def custom_matern_kernel(
X1: jnp.ndarray,
X2: jnp.ndarray,
params: Dict[str, jnp.ndarray],
noise: jnp.ndarray,
**kwargs: float
) -> jnp.ndarray:
"""
Matern 3/2 kernel with ARD support and numerical stability

Args:
    X1: Input array of shape (n_samples, n_features)
    X2: Input array of shape (n_samples, n_features)
    params: Dictionary of kernel parameters
        - k_length: lengthscale (array-like for ARD)
        - k_scale: output scale (scalar)
    noise: Observation noise variance
    **jitter: Numerical stability term
Returns:
    Covariance matrix of shape (n_samples, n_samples)
"""
# Get parameters with safety checks
lengthscale = params["k_length"]
variance = params["k_scale"] ** 2  # Square to match standard convention
jitter = kwargs.get("jitter", 1e-6)
# Reshape for single-feature datasets
if X1.ndim == 1:
    X1 = X1.reshape(-1, 1)
if X2.ndim == 1:
    X2 = X2.reshape(-1, 1)
      
# ARD handling
if lengthscale.ndim == 0 or lengthscale.shape[0] == 1:
    # Isotropic kernel
    X1_scaled = X1 / lengthscale
    X2_scaled = X2 / lengthscale
else:
    # ARD kernel
    X1_scaled = X1 / jnp.sqrt(lengthscale)
    X2_scaled = X2 / jnp.sqrt(lengthscale)
  
# Compute pairwise distances
D = jnp.sum(X1_scaled**2, axis=1)[:, None] + \
    jnp.sum(X2_scaled**2, axis=1) - \
    2 * jnp.dot(X1_scaled, X2_scaled.T)
D = jnp.sqrt(jnp.maximum(D, 1e-12))  # Ensure non-negative

# Matern 3/2 formula
sqrt3 = jnp.sqrt(3.0)
K = variance * (1 + sqrt3 * D) * jnp.exp(-sqrt3 * D)

# Add noise and jitter to diagonal only if X1 and X2 are the same and have the same values
# Removed jnp.allclose for static shape comparison
if X1.shape == X2.shape and X1.shape[0] > 0: # Check for non-empty arrays
    K = K.at[jnp.diag_indices(X1.shape[0])].add(noise + jitter) 
    
return K

def custom_matern_prior(input_dim: int) -> Callable[[], Dict[str, jnp.ndarray]]:
"""
Factory function returning prior distributions for Matern kernel parameters
with proper ARD handling
"""
def _prior() -> Dict[str, jnp.ndarray]:
# Lengthscale prior (ARD compatible)
with numpyro.plate("ard", input_dim):
k_length = numpyro.sample(
"k_length",
dist.LogNormal(0.0, 1.0).mask(False) # Allows gradient flow
)
# Output scale prior
k_scale = numpyro.sample(
"k_scale",
dist.LogNormal(0.0, 1.0).mask(False))

    return {"k_length": k_length, "k_scale": k_scale}

return _prior

But, having error as , "TypeError: custom_matern_kernel() missing 1 required positional argument: 'noise'1'" when implementing the following code:

Get random number generator keys (see JAX documentation for why it is neccessary)

rng_key, rng_key_predict = gpax.utils.get_keys()

Get random number generator keys for training and prediction

key1, key2 = gpax.utils.get_keys()

Initialize model

#gp_model = gpax.ExactGP(1, kernel=custom_matern_kernel,kernel_prior=custom_matern_prior(input_dim=1),noise_prior_dist=dist.HalfNormal(0.1))

Initialize GP with custom Matern kernel and priors

gp_model = gpax.ExactGP(
input_dim=1, # Example for 2D input
kernel=custom_matern_kernel,
kernel_prior=custom_matern_prior(input_dim=1), # Pass input_dim here
)

Run Hamiltonian Monte Carlo to obtain posterior samples for kernel parameters and model noise

gp_model.fit(key1, X, y, num_chains=1)
posterior_mean, f_samples = gp_model.predict(key2, X_test, n=200)

[ziatdinovmax]

Noise is supposed to be a keyword argument in custom kernel, set to 0 by default. So just set noise=0 in your custom_matern_kernel and then it should work:

def custom_matern_kernel(
    X1: jnp.ndarray,
    X2: jnp.ndarray,
    params: Dict[str, jnp.ndarray],
    noise: jnp.ndarray = 0,  # set noise to 0 by default
    **kwargs: float
    ) -> jnp.ndarray:

Btw, gpax already has a built-in support for ARD whenever the number of input dimensions is larger than 1.