An error occurred while executing the following cell:
------------------
# Predict using kernel ridge
X_plot = jnp.linspace(0, 20, 10000)[:, None]
stime = time.time()
y_kr = kr.predict(X_plot)
print("Time for KRR prediction: %.3f" % (time.time() - stime))

# Predict using gp.predict
X_plot = X_plot.reshape(
    -1,
)
y = y.reshape(
    -1,
)


def build_gp(theta_, X):
    kernel = kernels.ExpSineSquared(scale=jnp.exp(theta_["log_scale"]), gamma=jnp.exp(theta_["log_gamma"]))
    gp = GaussianProcess(kernel, X, diag=jnp.exp(theta_["log_diag"]))
    return gp


# predict without variance
stime = time.time()
gp = build_gp(soln.params, X)
y_mu = gp.predict(y, X_plot, return_var=False)
print("Time for GPR prediction: %.3f" % (time.time() - stime))

# predict with variance
stime = time.time()
gp = build_gp(soln.params, X)
y_mu, y_var = gp.predict(y, X_plot, return_var=True)
print("Time for GPR prediction with standard-deviation: %.3f" % (time.time() - stime))
------------------

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_6174/2032922369.py in <module>
     23 stime = time.time()
     24 gp = build_gp(soln.params, X)
---> 25 y_mu = gp.predict(y, X_plot, return_var=False)
     26 print("Time for GPR prediction: %.3f" % (time.time() - stime))
     27 

    [... skipping hidden 11 frame]

~/miniconda3/envs/py37/lib/python3.7/site-packages/tinygp/gp.py in predict(self, y, X_test, kernel, include_mean, return_var, return_cov)
    262         """
    263         _, cond = self.condition(
--> 264             y, X_test, kernel=kernel, include_mean=include_mean
    265         )
    266         if return_var:

~/miniconda3/envs/py37/lib/python3.7/site-packages/tinygp/gp.py in condition(self, y, X_test, diag, noise, include_mean, kernel)
    181             if not jax.tree_util.tree_reduce(lambda a, b: a and b, matches):
    182                 raise ValueError(
--> 183                     "`X_test` must have the same tree structure as the input `X`, "
    184                     "and all but the leading dimension must have matching sizes"
    185                 )

ValueError: `X_test` must have the same tree structure as the input `X`, and all but the leading dimension must have matching sizes
ValueError: `X_test` must have the same tree structure as the input `X`, and all but the leading dimension must have matching sizes