gpjax 0.10.2__py3-none-any.whl → 0.11.1__py3-none-any.whl

gpjax/__init__.py

@@ -32,14 +32,15 @@ from gpjax.citation import cite
 from gpjax.dataset import Dataset
 from gpjax.fit import (
     fit,
+    fit_lbfgs,
     fit_scipy,
 )
 
 __license__ = "MIT"
-__description__ = "Didactic Gaussian processes in JAX"
+__description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.10.2"
+__version__ = "0.11.1"
 
 __all__ = [
     "base",
@@ -56,5 +57,6 @@ __all__ = [
     "fit",
     "Module",
     "param_field",
+    "fit_lbfgs",
     "fit_scipy",
 ]

gpjax/distributions.py

@@ -15,77 +15,76 @@
 
 
 from beartype.typing import (
-    Any,
     Optional,
-    Tuple,
-    TypeVar,
 )
 import cola
+from cola.linalg.decompositions import Cholesky
 from cola.ops import (
-    Identity,
     LinearOperator,
 )
 from jax import vmap
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+from numpyro.distributions import constraints
+from numpyro.distributions.distribution import Distribution
+from numpyro.distributions.util import is_prng_key
 
 from gpjax.lower_cholesky import lower_cholesky
 from gpjax.typing import (
     Array,
-    KeyArray,
     ScalarFloat,
 )
 
-tfd = tfp.distributions
-
-from cola.linalg.decompositions import Cholesky
-
 
-class GaussianDistribution(tfd.Distribution):
-    r"""Multivariate Gaussian distribution with a linear operator scale matrix."""
-
-    # TODO: Consider `distrax.transformed.Transformed` object. Can we create a LinearOperator to `distrax.bijector` representation
-    # and modify `distrax.MultivariateNormalFromBijector`?
-    # TODO: Consider natural and expectation parameterisations in future work.
-    # TODO: we don't really need to inherit from `tfd.Distribution` here
+class GaussianDistribution(Distribution):
+    support = constraints.real_vector
 
     def __init__(
         self,
-        loc: Optional[Float[Array, " N"]] = None,
-        scale: Optional[LinearOperator] = None,
-    ) -> None:
-        r"""Initialises the distribution.
-
-        Args:
-            loc: the mean of the distribution as an array of shape (n_points,).
-            scale: the scale matrix of the distribution as a LinearOperator object.
-        """
-        _check_loc_scale(loc, scale)
+        loc: Optional[Float[Array, " N"]],
+        scale: Optional[LinearOperator],
+        validate_args=None,
+    ):
+        self.loc = loc
+        self.scale = cola.PSD(scale)
+        batch_shape = ()
+        event_shape = jnp.shape(self.loc)
+        super().__init__(batch_shape, event_shape, validate_args=validate_args)
 
-        # Find dimensionality of the distribution.
-        if loc is not None:
-            num_dims = loc.shape[-1]
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        # Obtain covariance root.
+        covariance_root = lower_cholesky(self.scale)
 
-        elif scale is not None:
-            num_dims = scale.shape[-1]
+        # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
+        white_noise = jr.normal(
+            key, shape=sample_shape + self.batch_shape + self.event_shape
+        )
 
-        # Set the location to zero vector if unspecified.
-        if loc is None:
-            loc = jnp.zeros((num_dims,))
+        # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
+        def affine_transformation(_x):
+            return self.loc + covariance_root @ _x
 
-        # If not specified, set the scale to the identity matrix.
-        if scale is None:
-            scale = Identity(shape=(num_dims, num_dims), dtype=loc.dtype)
-
-        self.loc = loc
-        self.scale = cola.PSD(scale)
+        return vmap(affine_transformation)(white_noise)
 
+    @property
     def mean(self) -> Float[Array, " N"]:
         r"""Calculates the mean."""
         return self.loc
 
+    @property
+    def variance(self) -> Float[Array, " N"]:
+        r"""Calculates the variance."""
+        return cola.diag(self.scale)
+
+    def entropy(self) -> ScalarFloat:
+        r"""Calculates the entropy of the distribution."""
+        return 0.5 * (
+            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi))
+            + cola.logdet(self.scale, Cholesky(), Cholesky())
+        )
+
     def median(self) -> Float[Array, " N"]:
         r"""Calculates the median."""
         return self.loc
@@ -98,25 +97,19 @@ class GaussianDistribution(tfd.Distribution):
         r"""Calculates the covariance matrix."""
         return self.scale.to_dense()
 
-    def variance(self) -> Float[Array, " N"]:
-        r"""Calculates the variance."""
-        return cola.diag(self.scale)
+    @property
+    def covariance_matrix(self) -> Float[Array, "N N"]:
+        r"""Calculates the covariance matrix."""
+        return self.covariance()
 
     def stddev(self) -> Float[Array, " N"]:
         r"""Calculates the standard deviation."""
         return jnp.sqrt(cola.diag(self.scale))
 
-    @property
-    def event_shape(self) -> Tuple:
-        r"""Returns the event shape."""
-        return self.loc.shape[-1:]
-
-    def entropy(self) -> ScalarFloat:
-        r"""Calculates the entropy of the distribution."""
-        return 0.5 * (
-            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi))
-            + cola.logdet(self.scale, Cholesky(), Cholesky())
-        )
+    #     @property
+    #     def event_shape(self) -> Tuple:
+    #         r"""Returns the event shape."""
+    #         return self.loc.shape[-1:]
 
     def log_prob(self, y: Float[Array, " N"]) -> ScalarFloat:
         r"""Calculates the log pdf of the multivariate Gaussian.
@@ -141,42 +134,39 @@ class GaussianDistribution(tfd.Distribution):
             + diff.T @ cola.solve(sigma, diff, Cholesky())
         )
 
-    def _sample_n(self, key: KeyArray, n: int) -> Float[Array, "n N"]:
-        r"""Samples from the distribution.
+    #     def _sample_n(self, key: KeyArray, n: int) -> Float[Array, "n N"]:
+    #         r"""Samples from the distribution.
 
-        Args:
-            key (KeyArray): The key to use for sampling.
+    #         Args:
+    #             key (KeyArray): The key to use for sampling.
 
-        Returns:
-            The samples as an array of shape (n_samples, n_points).
-        """
-        # Obtain covariance root.
-        sqrt = lower_cholesky(self.scale)
+    #         Returns:
+    #             The samples as an array of shape (n_samples, n_points).
+    #         """
+    #         # Obtain covariance root.
+    #         sqrt = lower_cholesky(self.scale)
 
-        # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
-        Z = jr.normal(key, shape=(n, *self.event_shape))
+    #         # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
+    #         Z = jr.normal(key, shape=(n, *self.event_shape))
 
-        # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
-        def affine_transformation(x):
-            return self.loc + sqrt @ x
+    #         # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
+    #         def affine_transformation(x):
+    #             return self.loc + sqrt @ x
 
-        return vmap(affine_transformation)(Z)
+    #         return vmap(affine_transformation)(Z)
 
-    def sample(
-        self, seed: KeyArray, sample_shape: Tuple[int, ...]
-    ):  # pylint: disable=useless-super-delegation
-        r"""See `Distribution.sample`."""
-        return self._sample_n(
-            seed, sample_shape[0]
-        )  # TODO this looks weird, why ignore the second entry?
+    #     def sample(
+    #         self, seed: KeyArray, sample_shape: Tuple[int, ...]
+    #     ):  # pylint: disable=useless-super-delegation
+    #         r"""See `Distribution.sample`."""
+    #         return self._sample_n(
+    #             seed, sample_shape[0]
+    #         )  # TODO this looks weird, why ignore the second entry?
 
     def kl_divergence(self, other: "GaussianDistribution") -> ScalarFloat:
         return _kl_divergence(self, other)
 
 
-DistrT = TypeVar("DistrT", bound=tfd.Distribution)
-
-
 def _check_and_return_dimension(
     q: GaussianDistribution, p: GaussianDistribution
 ) -> int:
@@ -245,37 +235,37 @@ def _kl_divergence(q: GaussianDistribution, p: GaussianDistribution) -> ScalarFl
     ) / 2.0
 
 
-def _check_loc_scale(loc: Optional[Any], scale: Optional[Any]) -> None:
-    r"""Checks that the inputs are correct."""
-    if loc is None and scale is None:
-        raise ValueError("At least one of `loc` or `scale` must be specified.")
-
-    if loc is not None and loc.ndim < 1:
-        raise ValueError("The parameter `loc` must have at least one dimension.")
-
-    if scale is not None and len(scale.shape) < 2:  # scale.ndim < 2:
-        raise ValueError(
-            "The `scale` must have at least two dimensions, but "
-            f"`scale.shape = {scale.shape}`."
-        )
-
-    if scale is not None and not isinstance(scale, LinearOperator):
-        raise ValueError(
-            f"The `scale` must be a CoLA LinearOperator but got {type(scale)}"
-        )
-
-    if scale is not None and (scale.shape[-1] != scale.shape[-2]):
-        raise ValueError(
-            f"The `scale` must be a square matrix, but `scale.shape = {scale.shape}`."
-        )
-
-    if loc is not None:
-        num_dims = loc.shape[-1]
-        if scale is not None and (scale.shape[-1] != num_dims):
-            raise ValueError(
-                f"Shapes are not compatible: `loc.shape = {loc.shape}` and "
-                f"`scale.shape = {scale.shape}`."
-            )
+# def _check_loc_scale(loc: Optional[Any], scale: Optional[Any]) -> None:
+#     r"""Checks that the inputs are correct."""
+#     if loc is None and scale is None:
+#         raise ValueError("At least one of `loc` or `scale` must be specified.")
+
+#     if loc is not None and loc.ndim < 1:
+#         raise ValueError("The parameter `loc` must have at least one dimension.")
+
+#     if scale is not None and len(scale.shape) < 2:  # scale.ndim < 2:
+#         raise ValueError(
+#             "The `scale` must have at least two dimensions, but "
+#             f"`scale.shape = {scale.shape}`."
+#         )
+
+#     if scale is not None and not isinstance(scale, LinearOperator):
+#         raise ValueError(
+#             f"The `scale` must be a CoLA LinearOperator but got {type(scale)}"
+#         )
+
+#     if scale is not None and (scale.shape[-1] != scale.shape[-2]):
+#         raise ValueError(
+#             f"The `scale` must be a square matrix, but `scale.shape = {scale.shape}`."
+#         )
+
+#     if loc is not None:
+#         num_dims = loc.shape[-1]
+#         if scale is not None and (scale.shape[-1] != num_dims):
+#             raise ValueError(
+#                 f"Shapes are not compatible: `loc.shape = {loc.shape}` and "
+#                 f"`scale.shape = {scale.shape}`."
+#             )
 
 
 __all__ = [

gpjax/fit.py

@@ -15,14 +15,14 @@
 
 import typing as tp
 
-from flax import nnx
 import jax
-from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
 import optax as ox
+from flax import nnx
+from jax.flatten_util import ravel_pytree
+from numpyro.distributions.transforms import Transform
 from scipy.optimize import minimize
-from tensorflow_probability.substrates.jax.bijectors import Bijector
 
 from gpjax.dataset import Dataset
 from gpjax.objectives import Objective
@@ -47,7 +47,7 @@ def fit(  # noqa: PLR0913
     objective: Objective,
     train_data: Dataset,
     optim: ox.GradientTransformation,
-    params_bijection: tp.Union[dict[Parameter, Bijector], None] = DEFAULT_BIJECTION,
+    params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
     key: KeyArray = jr.PRNGKey(42),
     num_iters: int = 100,
     batch_size: int = -1,
@@ -127,7 +127,6 @@ def fit(  # noqa: PLR0913
         _check_verbose(verbose)
 
     # Model state filtering
-
     graphdef, params, *static_state = nnx.split(model, Parameter, ...)
 
     # Parameters bijection to unconstrained space
@@ -253,6 +252,110 @@ def fit_scipy(  # noqa: PLR0913
     return model, history
 
 
+def fit_lbfgs(
+    *,
+    model: Model,
+    objective: Objective,
+    train_data: Dataset,
+    params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
+    max_iters: int = 100,
+    safe: bool = True,
+    max_linesearch_steps: int = 32,
+    gtol: float = 1e-5,
+) -> tuple[Model, jax.Array]:
+    r"""Train a Module model with respect to a supplied Objective function.
+
+    Uses Optax's LBFGS implementation and a jax.lax.while loop.
+
+     Args:
+         model: the model Module to be optimised.
+         objective: The objective function that we are optimising with
+             respect to.
+         train_data (Dataset): The training data to be used for the optimisation.
+         max_iters (int): The maximum number of optimisation steps to run. Defaults
+             to 500.
+         safe (bool): Whether to check the types of the inputs.
+         max_linesearch_steps (int): The maximum number of linesearch steps to use
+            for finding the stepsize.
+        gtol (float): Terminate the optimisation if the L2 norm of the gradient is
+            below this threshold.
+
+     Returns:
+         A tuple comprising the optimised model and final loss.
+    """
+    if safe:
+        # Check inputs
+        _check_model(model)
+        _check_train_data(train_data)
+        _check_num_iters(max_iters)
+
+    # Model state filtering
+    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+
+    # Parameters bijection to unconstrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection, inverse=True)
+
+    # Loss definition
+    def loss(params: nnx.State) -> ScalarFloat:
+        params = transform(params, params_bijection)
+        model = nnx.merge(graphdef, params, *static_state)
+        return objective(model, train_data)
+
+    # Initialise optimiser
+    optim = ox.lbfgs(
+        linesearch=ox.scale_by_zoom_linesearch(
+            max_linesearch_steps=max_linesearch_steps,
+            initial_guess_strategy="one",
+        )
+    )
+    opt_state = optim.init(params)
+    loss_value_and_grad = ox.value_and_grad_from_state(loss)
+
+    # Optimisation step.
+    def step(carry):
+        params, opt_state = carry
+
+        # Using optax's value_and_grad_from_state is more efficient given LBFGS uses a linesearch
+        # See https://optax.readthedocs.io/en/latest/api/utilities.html#optax.value_and_grad_from_state
+        loss_val, loss_gradient = loss_value_and_grad(params, state=opt_state)
+        updates, opt_state = optim.update(
+            loss_gradient,
+            opt_state,
+            params,
+            value=loss_val,
+            grad=loss_gradient,
+            value_fn=loss,
+        )
+        params = ox.apply_updates(params, updates)
+
+        return params, opt_state
+
+    def continue_fn(carry):
+        _, opt_state = carry
+        n = ox.tree_utils.tree_get(opt_state, "count")
+        g = ox.tree_utils.tree_get(opt_state, "grad")
+        g_l2_norm = ox.tree_utils.tree_l2_norm(g)
+        return (n == 0) | ((n < max_iters) & (g_l2_norm >= gtol))
+
+    # Optimisation loop
+    params, opt_state = jax.lax.while_loop(
+        continue_fn,
+        step,
+        (params, opt_state),
+    )
+    final_loss = ox.tree_utils.tree_get(opt_state, "value")
+
+    # Parameters bijection to constrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection)
+
+    # Reconstruct model
+    model = nnx.merge(graphdef, params, *static_state)
+
+    return model, final_loss
+
+
 def get_batch(train_data: Dataset, batch_size: int, key: KeyArray) -> Dataset:
     """Batch the data into mini-batches. Sampling is done with replacement.
 

gpjax/kernels/approximations/rff.py

@@ -68,7 +68,7 @@ class RFF(AbstractKernel):
 
             self.frequencies = Static(
                 self.base_kernel.spectral_density.sample(
-                    seed=key, sample_shape=(self.num_basis_fns, n_dims)
+                    key=key, sample_shape=(self.num_basis_fns, n_dims)
                 )
             )
         self.name = f"{self.base_kernel.name} (RFF)"

gpjax/kernels/nonstationary/arccosine.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -91,9 +94,9 @@ class ArcCosine(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         if isinstance(bias_variance, nnx.Variable):
             self.bias_variance = bias_variance

gpjax/kernels/nonstationary/linear.py

@@ -23,7 +23,7 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import NonNegativeReal
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -64,9 +64,9 @@ class Linear(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
     def __call__(
         self,

gpjax/kernels/nonstationary/polynomial.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -76,9 +79,9 @@ class Polynomial(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         self.name = f"Polynomial (degree {self.degree})"
 

gpjax/kernels/stationary/base.py

@@ -18,14 +18,17 @@ import beartype.typing as tp
 from flax import nnx
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.base import AbstractKernel
 from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -85,14 +88,14 @@ class StationaryKernel(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
 
             # static typing
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarFloat], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarFloat], self.variance)
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.Normal | npd.StudentT:
         r"""The spectral density of the kernel.
 
         Returns:

gpjax/kernels/stationary/matern12.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -48,5 +48,5 @@ class Matern12(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=1)

gpjax/kernels/stationary/matern32.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -54,5 +54,5 @@ class Matern32(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=3)

gpjax/kernels/stationary/matern52.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -53,5 +53,5 @@ class Matern52(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=5)

gpjax/kernels/stationary/rbf.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import squared_distance
@@ -44,5 +44,5 @@ class RBF(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfp.distributions.Normal:
-        return tfp.distributions.Normal(0.0, 1.0)
+    def spectral_density(self) -> npd.Normal:
+        return npd.Normal(0.0, 1.0)

gpjax/kernels/stationary/utils.py

@@ -14,17 +14,15 @@
 # ==============================================================================
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 
 from gpjax.typing import (
     Array,
     ScalarFloat,
 )
 
-tfd = tfp.distributions
 
-
-def build_student_t_distribution(nu: int) -> tfd.Distribution:
+def build_student_t_distribution(nu: int) -> npd.StudentT:
     r"""Build a Student's t distribution with a fixed smoothness parameter.
 
     For a fixed half-integer smoothness parameter, compute the spectral density of a
@@ -37,7 +35,7 @@ def build_student_t_distribution(nu: int) -> tfd.Distribution:
     -------
         tfp.Distribution: A Student's t distribution with the same smoothness parameter.
     """
-    dist = tfd.StudentT(df=nu, loc=0.0, scale=1.0)
+    dist = npd.StudentT(df=nu, loc=0.0, scale=1.0)
     return dist