gpjax 0.10.2__tar.gz → 0.11.0__tar.gz

{gpjax-0.10.2 → gpjax-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.10.2
+Version: 0.11.0
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -24,8 +24,8 @@ Requires-Dist: jax>=0.5.0
 Requires-Dist: jaxlib>=0.5.0
 Requires-Dist: jaxtyping>0.2.10
 Requires-Dist: numpy>=2.0.0
+Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
-Requires-Dist: tensorflow-probability>=0.24.0
 Requires-Dist: tqdm>4.66.2
 Description-Content-Type: text/markdown
 
@@ -138,65 +138,6 @@ jupytext --to notebook example.py
 jupytext --to py:percent example.ipynb
 ```
 
-# Simple example
-
-Let us import some dependencies and simulate a toy dataset $\mathcal{D}$.
-
-```python
-from jax import config
-
-config.update("jax_enable_x64", True)
-
-import gpjax as gpx
-from jax import grad, jit
-import jax.numpy as jnp
-import jax.random as jr
-import optax as ox
-
-key = jr.key(123)
-
-f = lambda x: 10 * jnp.sin(x)
-
-n = 50
-x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(n,1)).sort()
-y = f(x) + jr.normal(key, shape=(n,1))
-D = gpx.Dataset(X=x, y=y)
-
-# Construct the prior
-meanf = gpx.mean_functions.Zero()
-kernel = gpx.kernels.RBF()
-prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
-
-# Define a likelihood
-likelihood = gpx.likelihoods.Gaussian(num_datapoints = n)
-
-# Construct the posterior
-posterior = prior * likelihood
-
-# Define an optimiser
-optimiser = ox.adam(learning_rate=1e-2)
-
-# Obtain Type 2 MLEs of the hyperparameters
-opt_posterior, history = gpx.fit(
-    model=posterior,
-    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
-    train_data=D,
-    optim=optimiser,
-    num_iters=500,
-    safe=True,
-    key=key,
-)
-
-# Infer the predictive posterior distribution
-xtest = jnp.linspace(-3., 3., 100).reshape(-1, 1)
-latent_dist = opt_posterior(xtest, D)
-predictive_dist = opt_posterior.likelihood(latent_dist)
-
-# Obtain the predictive mean and standard deviation
-pred_mean = predictive_dist.mean()
-pred_std = predictive_dist.stddev()
-```
-
 # Installation
 
 ## Stable version

{gpjax-0.10.2 → gpjax-0.11.0}/README.md

@@ -107,65 +107,6 @@ jupytext --to notebook example.py
 jupytext --to py:percent example.ipynb
 ```
 
-# Simple example
-
-Let us import some dependencies and simulate a toy dataset $\mathcal{D}$.
-
-```python
-from jax import config
-
-config.update("jax_enable_x64", True)
-
-import gpjax as gpx
-from jax import grad, jit
-import jax.numpy as jnp
-import jax.random as jr
-import optax as ox
-
-key = jr.key(123)
-
-f = lambda x: 10 * jnp.sin(x)
-
-n = 50
-x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(n,1)).sort()
-y = f(x) + jr.normal(key, shape=(n,1))
-D = gpx.Dataset(X=x, y=y)
-
-# Construct the prior
-meanf = gpx.mean_functions.Zero()
-kernel = gpx.kernels.RBF()
-prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
-
-# Define a likelihood
-likelihood = gpx.likelihoods.Gaussian(num_datapoints = n)
-
-# Construct the posterior
-posterior = prior * likelihood
-
-# Define an optimiser
-optimiser = ox.adam(learning_rate=1e-2)
-
-# Obtain Type 2 MLEs of the hyperparameters
-opt_posterior, history = gpx.fit(
-    model=posterior,
-    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
-    train_data=D,
-    optim=optimiser,
-    num_iters=500,
-    safe=True,
-    key=key,
-)
-
-# Infer the predictive posterior distribution
-xtest = jnp.linspace(-3., 3., 100).reshape(-1, 1)
-latent_dist = opt_posterior(xtest, D)
-predictive_dist = opt_posterior.likelihood(latent_dist)
-
-# Obtain the predictive mean and standard deviation
-pred_mean = predictive_dist.mean()
-pred_std = predictive_dist.stddev()
-```
-
 # Installation
 
 ## Stable version

{gpjax-0.10.2 → gpjax-0.11.0}/docs/scripts/sharp_bits_figure.py

@@ -69,12 +69,12 @@ ax.set_xlim(-0.07, 0.25)
 plt.savefig("../_static/step_size_figure.png", bbox_inches="tight")
 
 # %%
-import tensorflow_probability.substrates.jax.bijectors as tfb
+import numpyro.distributions.transforms as npt
 
-bij = tfb.Exp()
+bij = npt.ExpTransform()
 
 x = np.linspace(0.05, 3.0, 6)
-y = np.asarray(bij.inverse(x))
+y = np.asarray(bij.inv(x))
 lval = 0.5
 rval = 0.52
 

{gpjax-0.10.2 → gpjax-0.11.0}/docs/sharp_bits.md

@@ -80,7 +80,7 @@ this value that we apply gradient updates to. When we wish to recover the constr
 value, we apply the inverse of the bijector, which is the exponential function in this
 case. This gives us back the blue cross.
 
-In GPJax, we supply bijective functions using [Tensorflow Probability](https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors).
+In GPJax, we supply bijective functions using [Numpyro](https://num.pyro.ai/en/stable/distributions.html#transforms).
 
 
 ## Positive-definiteness

{gpjax-0.10.2 → gpjax-0.11.0}/examples/barycentres.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -41,7 +41,7 @@ import jax.random as jr
 import jax.scipy.linalg as jsl
 from jaxtyping import install_import_hook
 import matplotlib.pyplot as plt
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from examples.utils import use_mpl_style
 
@@ -161,7 +161,7 @@ plt.show()
 
 
 # %%
-def fit_gp(x: jax.Array, y: jax.Array) -> tfd.MultivariateNormalFullCovariance:
+def fit_gp(x: jax.Array, y: jax.Array) -> npd.MultivariateNormal:
     if y.ndim == 1:
         y = y.reshape(-1, 1)
     D = gpx.Dataset(X=x, y=y)
@@ -204,9 +204,9 @@ def sqrtm(A: jax.Array):
 
 
 def wasserstein_barycentres(
-    distributions: tp.List[tfd.MultivariateNormalFullCovariance], weights: jax.Array
+    distributions: tp.List[npd.MultivariateNormal], weights: jax.Array
 ):
-    covariances = [d.covariance() for d in distributions]
+    covariances = [d.covariance_matrix for d in distributions]
     cov_stack = jnp.stack(covariances)
     stack_sqrt = jax.vmap(sqrtm)(cov_stack)
 
@@ -231,7 +231,7 @@ def wasserstein_barycentres(
 # %%
 weights = jnp.ones((n_datasets,)) / n_datasets
 
-means = jnp.stack([d.mean() for d in posterior_preds])
+means = jnp.stack([d.mean for d in posterior_preds])
 barycentre_mean = jnp.tensordot(weights, means, axes=1)
 
 step_fn = jax.jit(wasserstein_barycentres(posterior_preds, weights))
@@ -242,7 +242,7 @@ barycentre_covariance, sequence = jax.lax.scan(
 )
 L = jnp.linalg.cholesky(barycentre_covariance)
 
-barycentre_process = tfd.MultivariateNormalTriL(barycentre_mean, L)
+barycentre_process = npd.MultivariateNormal(barycentre_mean, scale_tril=L)
 
 # %% [markdown]
 # ## Plotting the result
@@ -254,7 +254,7 @@ barycentre_process = tfd.MultivariateNormalTriL(barycentre_mean, L)
 
 # %%
 def plot(
-    dist: tfd.MultivariateNormalTriL,
+    dist: npd.MultivariateNormal,
     ax,
     color: str,
     label: str = None,
@@ -262,8 +262,8 @@ def plot(
     linewidth: float = 1.0,
     zorder: int = 0,
 ):
-    mu = dist.mean()
-    sigma = dist.stddev()
+    mu = dist.mean
+    sigma = jnp.sqrt(dist.variance)
     ax.plot(xtest, mu, linewidth=linewidth, color=color, label=label, zorder=zorder)
     ax.fill_between(
         xtest.squeeze(),

{gpjax-0.10.2 → gpjax-0.11.0}/examples/classification.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -37,8 +37,8 @@ from jaxtyping import (
     install_import_hook,
 )
 import matplotlib.pyplot as plt
+import numpyro.distributions as npd
 import optax as ox
-import tensorflow_probability.substrates.jax as tfp
 
 from examples.utils import use_mpl_style
 from gpjax.lower_cholesky import lower_cholesky
@@ -50,7 +50,6 @@ with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
 
 
-tfd = tfp.distributions
 identity_matrix = jnp.eye
 
 # set the default style for plotting
@@ -120,7 +119,6 @@ print(type(posterior))
 # Optax's optimisers.
 
 # %%
-
 optimiser = ox.adam(learning_rate=0.01)
 
 opt_posterior, history = gpx.fit(
@@ -140,8 +138,8 @@ opt_posterior, history = gpx.fit(
 map_latent_dist = opt_posterior.predict(xtest, train_data=D)
 predictive_dist = opt_posterior.likelihood(map_latent_dist)
 
-predictive_mean = predictive_dist.mean()
-predictive_std = predictive_dist.stddev()
+predictive_mean = predictive_dist.mean
+predictive_std = jnp.sqrt(predictive_dist.variance)
 
 fig, ax = plt.subplots()
 ax.scatter(x, y, label="Observations", color=cols[0])
@@ -215,8 +213,6 @@ ax.legend()
 # datapoints below.
 
 # %%
-
-
 gram, cross_covariance = (kernel.gram, kernel.cross_covariance)
 jitter = 1e-6
 
@@ -246,7 +242,7 @@ L = jnp.linalg.cholesky(H + identity_matrix(D.n) * jitter)
 L_inv = jsp.linalg.solve_triangular(L, identity_matrix(D.n), lower=True)
 H_inv = jsp.linalg.solve_triangular(L.T, L_inv, lower=False)
 LH = jnp.linalg.cholesky(H_inv)
-laplace_approximation = tfd.MultivariateNormalTriL(f_hat.squeeze(), LH)
+laplace_approximation = npd.MultivariateNormal(f_hat.squeeze(), scale_tril=LH)
 
 
 # %% [markdown]
@@ -265,7 +261,7 @@ laplace_approximation = tfd.MultivariateNormalTriL(f_hat.squeeze(), LH)
 
 
 # %%
-def construct_laplace(test_inputs: Float[Array, "N D"]) -> tfd.MultivariateNormalTriL:
+def construct_laplace(test_inputs: Float[Array, "N D"]) -> npd.MultivariateNormal:
     map_latent_dist = opt_posterior.predict(xtest, train_data=D)
 
     Kxt = opt_posterior.prior.kernel.cross_covariance(x, test_inputs)
@@ -279,10 +275,10 @@ def construct_laplace(test_inputs: Float[Array, "N D"]) -> tfd.MultivariateNorma
     # Ktx Kxx⁻¹[ H⁻¹ ] Kxx⁻¹ Kxt
     laplace_cov_term = jnp.matmul(jnp.matmul(Kxx_inv_Kxt.T, H_inv), Kxx_inv_Kxt)
 
-    mean = map_latent_dist.mean()
-    covariance = map_latent_dist.covariance() + laplace_cov_term
+    mean = map_latent_dist.mean
+    covariance = map_latent_dist.covariance_matrix + laplace_cov_term
     L = jnp.linalg.cholesky(covariance)
-    return tfd.MultivariateNormalTriL(jnp.atleast_1d(mean.squeeze()), L)
+    return npd.MultivariateNormal(jnp.atleast_1d(mean.squeeze()), scale_tril=L)
 
 
 # %% [markdown]
@@ -291,8 +287,8 @@ def construct_laplace(test_inputs: Float[Array, "N D"]) -> tfd.MultivariateNorma
 laplace_latent_dist = construct_laplace(xtest)
 predictive_dist = opt_posterior.likelihood(laplace_latent_dist)
 
-predictive_mean = predictive_dist.mean()
-predictive_std = predictive_dist.stddev()
+predictive_mean = predictive_dist.mean
+predictive_std = jnp.sqrt(predictive_dist.variance)
 
 fig, ax = plt.subplots()
 ax.scatter(x, y, label="Observations", color=cols[0])

{gpjax-0.10.2 → gpjax-0.11.0}/examples/collapsed_vi.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax_beartype
 #     language: python
@@ -161,10 +161,10 @@ predictive_dist = opt_posterior.posterior.likelihood(latent_dist)
 
 inducing_points = opt_posterior.inducing_inputs.value
 
-samples = latent_dist.sample(seed=key, sample_shape=(20,))
+samples = latent_dist.sample(key=key, sample_shape=(20,))
 
-predictive_mean = predictive_dist.mean()
-predictive_std = predictive_dist.stddev()
+predictive_mean = predictive_dist.mean
+predictive_std = jnp.sqrt(predictive_dist.variance)
 
 fig, ax = plt.subplots()
 

{gpjax-0.10.2 → gpjax-0.11.0}/examples/constructing_new_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -24,6 +24,7 @@
 # %%
 # Enable Float64 for more stable matrix inversions.
 from jax import config
+from jax.nn import softplus
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import (
@@ -32,7 +33,9 @@ from jaxtyping import (
     install_import_hook,
 )
 import matplotlib.pyplot as plt
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
+from numpyro.distributions import constraints
+import numpyro.distributions.transforms as npt
 
 from examples.utils import use_mpl_style
 from gpjax.kernels.computations import DenseKernelComputation
@@ -225,9 +228,27 @@ def angular_distance(x, y, c):
     return jnp.abs((x - y + c) % (c * 2) - c)
 
 
-bij = tfb.SoftClip(low=jnp.array(4.0, dtype=jnp.float64))
+class ShiftedSoftplusTransform(npt.ParameterFreeTransform):
+    r"""
+    Transform from unconstrained space to the domain [4, infinity) via
+    :math:`y = 4 + \log(1 + \exp(x))`. The inverse is computed as
+    :math:`x = \log(\exp(y - 4) - 1)`.
+    """
 
-DEFAULT_BIJECTION["polar"] = bij
+    domain = constraints.real
+    codomain = constraints.interval(4.0, jnp.inf)  # updated codomain
+
+    def __call__(self, x):
+        return 4.0 + softplus(x)  # shift the softplus output by 4
+
+    def _inverse(self, y):
+        return npt._softplus_inv(y - 4.0)  # subtract the shift in the inverse
+
+    def log_abs_det_jacobian(self, x, y, intermediates=None):
+        return -softplus(-x)
+
+
+DEFAULT_BIJECTION["polar"] = ShiftedSoftplusTransform()
 
 
 class Polar(gpx.kernels.AbstractKernel):
@@ -307,7 +328,7 @@ opt_posterior, history = gpx.fit_scipy(
 
 # %%
 posterior_rv = opt_posterior.likelihood(opt_posterior.predict(angles, train_data=D))
-mu = posterior_rv.mean()
+mu = posterior_rv.mean
 one_sigma = posterior_rv.stddev()
 
 # %%

{gpjax-0.10.2 → gpjax-0.11.0}/examples/deep_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -238,8 +238,8 @@ opt_posterior, history = gpx.fit(
 latent_dist = opt_posterior(xtest, train_data=D)
 predictive_dist = opt_posterior.likelihood(latent_dist)
 
-predictive_mean = predictive_dist.mean()
-predictive_std = predictive_dist.stddev()
+predictive_mean = predictive_dist.mean
+predictive_std = jnp.sqrt(predictive_dist.variance)
 
 fig, ax = plt.subplots()
 ax.plot(x, y, "o", label="Observations", color=cols[0])

{gpjax-0.10.2 → gpjax-0.11.0}/examples/graph_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -124,7 +124,7 @@ true_kernel = gpx.kernels.GraphKernel(
 prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=true_kernel)
 
 fx = prior(x)
-y = fx.sample(seed=key, sample_shape=(1,)).reshape(-1, 1)
+y = fx.sample(key=key, sample_shape=(1,)).reshape(-1, 1)
 
 D = gpx.Dataset(X=x, y=y)
 
@@ -194,8 +194,8 @@ opt_posterior, training_history = gpx.fit_scipy(
 initial_dist = likelihood(posterior(x, D))
 predictive_dist = opt_posterior.likelihood(opt_posterior(x, D))
 
-initial_mean = initial_dist.mean()
-learned_mean = predictive_dist.mean()
+initial_mean = initial_dist.mean
+learned_mean = predictive_dist.mean
 
 rmse = lambda ytrue, ypred: jnp.sum(jnp.sqrt(jnp.square(ytrue - ypred)))
 

{gpjax-0.10.2 → gpjax-0.11.0}/examples/intro_to_gps.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.6
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -127,9 +127,9 @@ import jax.numpy as jnp
 import jax.random as jr
 import matplotlib as mpl
 import matplotlib.pyplot as plt
+import numpyro.distributions as npd
 import pandas as pd
 import seaborn as sns
-import tensorflow_probability.substrates.jax as tfp
 
 from examples.utils import (
     confidence_ellipse,
@@ -143,11 +143,10 @@ key = jr.key(42)
 
 
 cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
-tfd = tfp.distributions
 
-ud1 = tfd.Normal(0.0, 1.0)
-ud2 = tfd.Normal(-1.0, 0.5)
-ud3 = tfd.Normal(0.25, 1.5)
+ud1 = npd.Normal(0.0, 1.0)
+ud2 = npd.Normal(-1.0, 0.5)
+ud3 = npd.Normal(0.25, 1.5)
 
 xs = jnp.linspace(-5.0, 5.0, 500)
 
@@ -156,7 +155,7 @@ for d in [ud1, ud2, ud3]:
     ax.plot(
         xs,
         jnp.exp(d.log_prob(xs)),
-        label=f"$\\mathcal{{N}}({{{float(d.mean())}}},\\  {{{float(d.stddev())}}}^2)$",
+        label=f"$\\mathcal{{N}}({{{float(d.mean)}}},\\  {{{float(jnp.sqrt(d.variance))}}}^2)$",
     )
     ax.fill_between(xs, jnp.zeros_like(xs), jnp.exp(d.log_prob(xs)), alpha=0.2)
 ax.legend(loc="best")
@@ -190,12 +189,12 @@ ax.legend(loc="best")
 # %%
 key = jr.key(123)
 
-d1 = tfd.MultivariateNormalDiag(loc=jnp.zeros(2), scale_diag=jnp.ones(2))
-d2 = tfd.MultivariateNormalTriL(
-    jnp.zeros(2), jnp.linalg.cholesky(jnp.array([[1.0, 0.9], [0.9, 1.0]]))
+d1 = npd.MultivariateNormal(loc=jnp.zeros(2), covariance_matrix=jnp.diag(jnp.ones(2)))
+d2 = npd.MultivariateNormal(
+    jnp.zeros(2), scale_tril=jnp.linalg.cholesky(jnp.array([[1.0, 0.9], [0.9, 1.0]]))
 )
-d3 = tfd.MultivariateNormalTriL(
-    jnp.zeros(2), jnp.linalg.cholesky(jnp.array([[1.0, -0.5], [-0.5, 1.0]]))
+d3 = npd.MultivariateNormal(
+    jnp.zeros(2), scale_tril=jnp.linalg.cholesky(jnp.array([[1.0, -0.5], [-0.5, 1.0]]))
 )
 
 dists = [d1, d2, d3]
@@ -215,13 +214,21 @@ titles = [r"$\rho = 0$", r"$\rho = 0.9$", r"$\rho = -0.5$"]
 cmap = mpl.colors.LinearSegmentedColormap.from_list("custom", ["white", cols[1]], N=256)
 
 for a, t, d in zip([ax0, ax1, ax2], titles, dists, strict=False):
-    d_prob = d.prob(jnp.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)])).reshape(
-        xx.shape
+    d_prob = jnp.exp(
+        d.log_prob(jnp.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)]))
+    ).reshape(xx.shape)
+    cntf = a.contourf(
+        xx,
+        yy,
+        jnp.exp(d_prob),
+        levels=20,
+        antialiased=True,
+        cmap=cmap,
+        edgecolor="face",
     )
-    cntf = a.contourf(xx, yy, jnp.exp(d_prob), levels=20, antialiased=True, cmap=cmap, edgecolor="face")
     a.set_xlim(-2.75, 2.75)
     a.set_ylim(-2.75, 2.75)
-    samples = d.sample(seed=key, sample_shape=(5000,))
+    samples = d.sample(key=key, sample_shape=(5000,))
     xsample, ysample = samples[:, 0], samples[:, 1]
     confidence_ellipse(
         xsample, ysample, a, edgecolor="#3f3f3f", n_std=1.0, linestyle="--", alpha=0.8
@@ -274,13 +281,13 @@ for a, t, d in zip([ax0, ax1, ax2], titles, dists, strict=False):
 
 # %%
 n = 1000
-x = tfd.Normal(loc=0.0, scale=1.0).sample(seed=key, sample_shape=(n,))
+x = npd.Normal(loc=0.0, scale=1.0).sample(key, sample_shape=(n,))
 key, subkey = jr.split(key)
-y = tfd.Normal(loc=0.25, scale=0.5).sample(seed=subkey, sample_shape=(n,))
+y = npd.Normal(loc=0.25, scale=0.5).sample(subkey, sample_shape=(n,))
 key, subkey = jr.split(subkey)
-xfull = tfd.Normal(loc=0.0, scale=1.0).sample(seed=subkey, sample_shape=(n * 10,))
+xfull = npd.Normal(loc=0.0, scale=1.0).sample(subkey, sample_shape=(n * 10,))
 key, subkey = jr.split(subkey)
-yfull = tfd.Normal(loc=0.25, scale=0.5).sample(seed=subkey, sample_shape=(n * 10,))
+yfull = npd.Normal(loc=0.25, scale=0.5).sample(subkey, sample_shape=(n * 10,))
 key, subkey = jr.split(subkey)
 df = pd.DataFrame({"x": x, "y": y, "idx": jnp.ones(n)})