Saving and loading GP model using trace data gives wrong output shape

Questions version agnostic
1

Hello, I am trying to save and load the trace from a gp model to use the model for production. Here is my trial, everything seems to work except that I expect the

gp.predict to give me one output for each Xnew input, however it gives me an output shape equals the number of training samples.

def Train_for_production(collected_x_path,collected_y_path, trace_path, model_path, mcmc_samples_final=4000, mcmc_warmup_final=2000):

  #load inputs
  inputs_df=pd.read_csv(collected_x_path)
  #load outputs
  output_df=pd.read_csv(collected_y_path)
  #get n_features
  n_features=len(inputs_df.columns)-1

  #Normalize xs
  sc=StandardScaler()
  normalized_collected_xs=sc.fit_transform(np.array(inputs_df[inputs_df.columns[1:]])) 

  #Normalize ys
  y_scaler=StandardScaler()
  ys_normalized=y_scaler.fit_transform(np.array(output_df[output_df.columns[-1]]).reshape(-1,1))

  #Train the final surrogate on all the collected data 
  with pm.Model() as surrogate_model:

    pm.Data("X", normalized_collected_xs)

    #surrogate priors for each input feature (can be used to inform sensitivity analysis)
    features_priors=[]
    for i in range(n_features):
      features_priors.append(pm.HalfNormal(f"rho_{i+1}", 100))

    cov= pm.gp.cov.Matern32(n_features ,features_priors) #Squared Exponential ExpQuad
    mean_function=pm.gp.mean.Zero()
    gp = pm.gp.Marginal(mean_func=mean_function, cov_func=cov)


    y_ = gp.marginal_likelihood("y", X=normalized_collected_xs,
                                y=ys_normalized.reshape(ys_normalized.shape[0]), sigma=1e-12)
    with surrogate_model:
        marginal_post = pm.sample(mcmc_samples_final, tune=mcmc_warmup_final, nuts_sampler="numpyro", chains=4, return_inferencedata=True) #mcmc

  #Create InferenceData with BOTH trace AND training data
  
  idata = az.convert_to_inference_data(marginal_post)
  
  #Save everything
  az.to_netcdf(idata, trace_path)  # Now contains trace + training data

  # #Save the trace (posterior samples)
  # az.to_netcdf(marginal_post, trace_path)


  #Save the model
  pm.model_to_graphviz(surrogate_model).save(model_path)  
  #Save inputs scaler
  with open(inputs_scaler_path,'wb') as f:
    pickle.dump(sc, f)
  #Save output scaler
  with open(output_scaler_path,'wb') as f:
    pickle.dump(y_scaler, f)

this is the function I use for training and save the model on the dataset I have collected.

def production_surrogate_model(input_array, collected_x_path,collected_y_path, trace_path, model_path, inputs_scaler_path, output_scaler_path, n_chain_samples=500, predict_only_means=True, mcmc_samples=1000):
  
  #process the inputs 
  with open(inputs_scaler_path,'rb') as f:
    inputs_scaler = pickle.load(f)
  with open(output_scaler_path,'rb') as f:
    output_scaler = pickle.load(f)

  #load inputs
  inputs_df=pd.read_csv(collected_x_path)
  #load outputs
  output_df=pd.read_csv(collected_y_path)
  #get n_features
  n_features=len(inputs_df.columns)-1

  #Training data, neded for conditioning in the GP
  X_train=inputs_scaler.transform(np.array(inputs_df[inputs_df.columns[1:]])) 
  y_train=output_scaler.transform(np.array(output_df[output_df.columns[-1]]).reshape(-1,1))

  #input data for prediction
  normalized_inputs=inputs_scaler.transform(input_array)

  print(normalized_inputs.shape)

  #reload the trace
  trace = az.from_netcdf(trace_path)

  with pm.Model() as surrogate_model:

    features_priors=[]
    for i in range(n_features):
      features_priors.append(pm.HalfNormal(f"rho_{i+1}", 100))

    cov= pm.gp.cov.Matern32(n_features ,features_priors) #Squared Exponential ExpQuad
    mean_function=pm.gp.mean.Zero()

    gp = pm.gp.Marginal(mean_func=mean_function, cov_func=cov)
    #gp.marginal_likelihood("y", X=X_train, y=y_train, sigma=1e-12)
    y_ = gp.marginal_likelihood("y", X=X_train, y=y_train, sigma=1e-12)

  if predict_only_means:
    with surrogate_model:
      parmeters_dict={f"rho_{f_indx+1}":np.array(az.InferenceData.mean(trace).posterior[[f"rho_{f_indx+1}"]].to_array())[0] for f_indx in range(n_features)}
      pred, cov_diag = gp.predict(Xnew=normalized_inputs, point=parmeters_dict, diag=True)
    pred=output_scaler.inverse_transform(pred.reshape(-1,1))
    return pred
  
  else:
    #get full posterior distribution
    with surrogate_model:
      f_pred = gp.conditional('f_pred', normalized_inputs)
      
    with surrogate_model:
      pred_samples = pm.sample_posterior_predictive(trace=trace.sel(draw=slice(0, mcmc_samples)) ,var_names=["f_pred"], extend_inferencedata=True, random_seed=710)
    f_pred_samples = az.extract(pred_samples, group="posterior_predictive", var_names=["f_pred"])
    
    f_pred_samples =np.array(f_pred_samples)
    f_pred_samples=output_scaler.inverse_transform(f_pred_samples.reshape(-1,1))
    print(f_pred_samples.shape)
    return f_pred_samples

and this is the one for loading the model and using it for prediction.

Thank you!