From 193460d053674aa578712bd7fcc0941f3ad0c13e Mon Sep 17 00:00:00 2001
From: xImoZA <alekcandra.ri@gmail.com>
Date: Tue, 25 Nov 2025 14:47:47 +0300
Subject: [PATCH] refactor: distribution methods output types

---
 rework_pysatl_mpest/core/mixture.py           |  22 +--
 rework_pysatl_mpest/distributions/beta.py     |  66 ++++++--
 rework_pysatl_mpest/distributions/cauchy.py   |  46 ++++--
 .../distributions/continuous_dist.py          |  16 +-
 .../distributions/exponential.py              |  45 ++++--
 rework_pysatl_mpest/distributions/normal.py   |  22 ++-
 rework_pysatl_mpest/distributions/pareto.py   |  61 ++++++--
 rework_pysatl_mpest/distributions/uniform.py  |  74 +++++++--
 rework_pysatl_mpest/distributions/weibull.py  |  75 ++++++---
 rework_tests/unit/core/test_mixture.py        |  74 ++++++---
 rework_tests/unit/distributions/test_beta.py  | 144 +++++++++++++++---
 .../unit/distributions/test_cauchy.py         | 106 +++++++++++--
 .../unit/distributions/test_exponential.py    |  83 ++++++++--
 .../unit/distributions/test_normal.py         |  91 +++++++++--
 .../unit/distributions/test_pareto.py         |  84 ++++++++--
 .../unit/distributions/test_uniform.py        | 131 ++++++++++++++--
 .../unit/distributions/test_weibull.py        |  99 ++++++++++--
 17 files changed, 1032 insertions(+), 207 deletions(-)

diff --git a/rework_pysatl_mpest/core/mixture.py b/rework_pysatl_mpest/core/mixture.py
index 4c6789f..dd23af7 100644
--- a/rework_pysatl_mpest/core/mixture.py
+++ b/rework_pysatl_mpest/core/mixture.py
@@ -255,7 +255,7 @@ def remove_component(self, component_idx: int):
         self._cached_weights = None
         self._sorted_pairs_cache = None
 
-    def pdf(self, X: ArrayLike) -> NDArray[DType]:
+    def pdf(self, X: ArrayLike) -> DType | NDArray[DType]:
         """Probability Density Function of the mixture.
 
         The PDF is computed as the weighted sum of the PDFs of its
@@ -268,15 +268,16 @@ def pdf(self, X: ArrayLike) -> NDArray[DType]:
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
         component_pdfs = np.array([comp.pdf(X) for comp in self.components])
-        return np.asarray(np.dot(self.weights, component_pdfs))
+        return np.dot(self.weights, component_pdfs)
 
-    def lpdf(self, X: ArrayLike) -> NDArray[DType]:
+    def lpdf(self, X: ArrayLike) -> DType | NDArray[DType]:
         """Logarithms of the Probability Density Function.
 
         Parameters
@@ -286,15 +287,18 @@ def lpdf(self, X: ArrayLike) -> NDArray[DType]:
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
-        X = np.atleast_1d(X).astype(self.dtype)
+        X = np.asarray(X, dtype=self.dtype)
         component_lpdfs = np.array([comp.lpdf(X) for comp in self.components])
-        log_weights = self.log_weights
-        log_terms = log_weights[:, np.newaxis] + component_lpdfs
-        return logsumexp(log_terms, axis=0)  # type: ignore
+        broadcast_shape = (self.n_components,) + (1,) * X.ndim
+        log_weights = self.log_weights.reshape(broadcast_shape)
+        log_terms = log_weights + component_lpdfs
+
+        return logsumexp(log_terms, axis=0)
 
     def loglikelihood(self, X: ArrayLike) -> DType:
         """Log-likelihood of the complete data :attr:`X`.
diff --git a/rework_pysatl_mpest/distributions/beta.py b/rework_pysatl_mpest/distributions/beta.py
index c71eee4..32b45c3 100644
--- a/rework_pysatl_mpest/distributions/beta.py
+++ b/rework_pysatl_mpest/distributions/beta.py
@@ -109,10 +109,12 @@ def pdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
 
         """
+
         X = np.asarray(X, dtype=self.dtype)
         return np.exp(self.lpdf(X))
 
@@ -137,13 +139,16 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(
+        result = np.where(
             (P >= 0) & (P <= 1),
             (
                 self.left_border
@@ -152,6 +157,10 @@ def ppf(self, P):
             dtype(np.nan),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
 
@@ -177,8 +186,9 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
@@ -189,7 +199,7 @@ def lpdf(self, X):
         log_pdf_standard = beta_dist.logpdf(Z, self.alpha, self.beta).astype(dtype)
         result = log_pdf_standard - np.log(self.right_border - self.left_border)
 
-        return np.atleast_1d(result)
+        return result
 
     def _dlog_alpha(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`alpha` parameter.
@@ -214,15 +224,17 @@ def _dlog_alpha(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`alpha` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_bounds = (self.left_border < X) & (self.right_border >= X)
-        return np.where(
+        result = np.where(
             in_bounds,
             np.log(X - self.left_border)
             - np.log(self.right_border - self.left_border)
@@ -230,6 +242,10 @@ def _dlog_alpha(self, X):
             dtype(0.0),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def _dlog_beta(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`beta` parameter.
 
@@ -253,15 +269,17 @@ def _dlog_beta(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`beta` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_bounds = (self.left_border < X) & (self.right_border >= X)
-        return np.where(
+        result = np.where(
             in_bounds,
             np.log(self.right_border - X)
             - np.log(self.right_border - self.left_border)
@@ -269,6 +287,10 @@ def _dlog_beta(self, X):
             dtype(0.0),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def _dlog_left_border(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`left_border` parameter.
 
@@ -290,15 +312,17 @@ def _dlog_left_border(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`left_border` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_bounds = (self.left_border < X) & (self.right_border >= X)
-        return np.where(
+        result = np.where(
             in_bounds,
             (
                 ((self.alpha + self.beta - dtype(1)) / (self.right_border - self.left_border))
@@ -307,6 +331,10 @@ def _dlog_left_border(self, X):
             dtype(0.0),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def _dlog_right_border(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`right_border` parameter.
 
@@ -328,14 +356,17 @@ def _dlog_right_border(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`right_border` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_bounds = (self.left_border < X) & (self.right_border >= X)
-        return np.where(
+        result = np.where(
             in_bounds,
             (
                 ((self.beta - dtype(1)) / (self.right_border - X))
@@ -344,6 +375,10 @@ def _dlog_right_border(self, X):
             dtype(0.0),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
 
@@ -361,7 +396,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -378,6 +416,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/cauchy.py b/rework_pysatl_mpest/distributions/cauchy.py
index 3e5b486..aad4c09 100644
--- a/rework_pysatl_mpest/distributions/cauchy.py
+++ b/rework_pysatl_mpest/distributions/cauchy.py
@@ -82,14 +82,13 @@ def pdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
-
-        return dtype(1.0) / (dtype(np.pi) * self.scale * (dtype(1.0) + ((X - self.loc) / self.scale) ** 2))
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -110,13 +109,16 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(
+        result = np.where(
             (P >= 0) & (P <= 1),
             np.where(
                 (P == 0) | (P == 1),
@@ -125,6 +127,9 @@ def ppf(self, P):
             ),
             dtype(np.nan),
         )
+        if is_scalar:
+            return result[()]
+        return result
 
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
@@ -145,8 +150,9 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
@@ -179,14 +185,20 @@ def _dlog_loc(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`loc` for each point in ::attr`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return (dtype(2) * X - dtype(2) * self.loc) / (self.scale**2 + X**2 - dtype(2) * self.loc * X + self.loc**2)
+        result = (dtype(2) * X - dtype(2) * self.loc) / (self.scale**2 + X**2 - dtype(2) * self.loc * X + self.loc**2)
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_scale(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`scale` parameter.
@@ -208,16 +220,23 @@ def _dlog_scale(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`rate` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return (-(self.scale**2) + X**2 - dtype(2) * self.loc * X + self.loc**2) / (
+        result = (-(self.scale**2) + X**2 - dtype(2) * self.loc * X + self.loc**2) / (
             self.scale**3 + self.scale * (X**2) - dtype(2) * self.loc * self.scale * X + self.scale * self.loc**2
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
 
@@ -235,7 +254,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -250,6 +272,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/continuous_dist.py b/rework_pysatl_mpest/distributions/continuous_dist.py
index 58aca34..515932e 100644
--- a/rework_pysatl_mpest/distributions/continuous_dist.py
+++ b/rework_pysatl_mpest/distributions/continuous_dist.py
@@ -202,6 +202,7 @@ def set_params_from_vector(self, param_names: Sequence[str], vector: Sequence[Un
     @property
     def dtype(self) -> type[DType]:
         """type[DType]: The numpy data type of the distribution's outputs."""
+
         return self._dtype
 
     @property
@@ -221,7 +222,7 @@ def params_to_optimize(self) -> set[str]:
         return self.params - self._fixed_params
 
     @abstractmethod
-    def pdf(self, X: ArrayLike) -> NDArray[DType]:
+    def pdf(self, X: ArrayLike) -> DType | NDArray[DType]:
         """Probability Density Function.
 
         Parameters
@@ -231,8 +232,9 @@ def pdf(self, X: ArrayLike) -> NDArray[DType]:
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
     @abstractmethod
@@ -249,8 +251,9 @@ def ppf(self, P: ArrayLike) -> NDArray[DType]:
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
     @abstractmethod
@@ -268,8 +271,9 @@ def lpdf(self, X: ArrayLike) -> NDArray[DType]:
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
     @abstractmethod
@@ -325,6 +329,7 @@ def astype(self, new_dtype: type[DType]) -> "ContinuousDistribution[DType]":
             specified `new_dtype`, or the original instance if the `dtype` is
             unchanged.
         """
+
         if self._dtype is new_dtype:
             return self
 
@@ -343,6 +348,7 @@ def __copy__(self) -> "ContinuousDistribution[DType]":
         ContinuousDistribution[DType]
             A new instance of the distribution, identical to the original.
         """
+
         params_dict = {p: getattr(self, p) for p in self.params}
 
         new_instance = self.__class__(**params_dict, dtype=self.dtype)
@@ -366,6 +372,7 @@ def __eq__(self, other: object):
         bool
             True if the distributions are equal, False otherwise.
         """
+
         if not isinstance(other, ContinuousDistribution):
             return NotImplemented
 
@@ -392,6 +399,7 @@ def __hash__(self) -> int:
         int
             The hash value of the distribution object.
         """
+
         sorted_params = sorted(list(self.params))
         param_values = tuple(self.get_params_vector(sorted_params))
 
diff --git a/rework_pysatl_mpest/distributions/exponential.py b/rework_pysatl_mpest/distributions/exponential.py
index 772ba00..a549a26 100644
--- a/rework_pysatl_mpest/distributions/exponential.py
+++ b/rework_pysatl_mpest/distributions/exponential.py
@@ -87,9 +87,7 @@ def pdf(self, X):
         """
 
         X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
-
-        return np.where(self.loc <= X, self.rate * np.exp(-self.rate * (X - self.loc)), dtype(0.0))
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -107,14 +105,19 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where((P >= 0) & (P <= 1), self.loc - np.log(dtype(1) - P) / self.rate, dtype(np.nan))
+        result = np.where((P >= 0) & (P <= 1), self.loc - np.log(dtype(1) - P) / self.rate, dtype(np.nan))
+        if is_scalar:
+            return result[()]
+        return result
 
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
@@ -132,14 +135,19 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(self.loc <= X, np.log(self.rate) - self.rate * (X - self.loc), dtype(-np.inf))
+        result = np.where(self.loc <= X, np.log(self.rate) - self.rate * (X - self.loc), dtype(-np.inf))
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_loc(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`loc` parameter.
@@ -159,14 +167,19 @@ def _dlog_loc(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`loc` for each point in ::attr`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(self.loc <= X, self.rate, dtype(0.0))
+        result = np.where(self.loc <= X, self.rate, dtype(0.0))
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_rate(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`rate` parameter.
@@ -186,13 +199,19 @@ def _dlog_rate(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`rate` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
-        return np.where(self.loc <= X, dtype(1.0) / self.rate - (X - self.loc), dtype(0.0))
+
+        result = np.where(self.loc <= X, dtype(1.0) / self.rate - (X - self.loc), dtype(0.0))
+        if is_scalar:
+            return result[()]
+        return result
 
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
@@ -211,8 +230,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -227,6 +248,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/normal.py b/rework_pysatl_mpest/distributions/normal.py
index 30da38d..54e9af3 100644
--- a/rework_pysatl_mpest/distributions/normal.py
+++ b/rework_pysatl_mpest/distributions/normal.py
@@ -83,15 +83,13 @@ def pdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
-
-        z = (X - self.loc) / self.scale
-        return np.exp(-(z**2) / dtype(2.0)) / (self.scale * np.sqrt(dtype(2.0) * dtype(np.pi)))
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -107,13 +105,17 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         result = norm.ppf(P, loc=self.loc, scale=self.scale)
 
+        if is_scalar:
+            return self.dtype(result)
         return np.asarray(result, dtype=self.dtype)
 
     def lpdf(self, X):
@@ -133,8 +135,9 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
@@ -173,8 +176,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -188,6 +193,9 @@ def log_gradients(self, X):
             return np.empty((len(X), 0), dtype=self.dtype)
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
+
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/pareto.py b/rework_pysatl_mpest/distributions/pareto.py
index 2c9c696..03cbb6d 100644
--- a/rework_pysatl_mpest/distributions/pareto.py
+++ b/rework_pysatl_mpest/distributions/pareto.py
@@ -77,13 +77,21 @@ def pdf(self, X):
 
         where :math:`\\alpha` is the :attr:`shape` parameter and :math:`\\beta` is the
         :attr:`scale` parameter. The function is zero for :math:`x < \\beta`.
+
+        Parameters
+        ----------
+        X : ArrayLike
+            The input data points at which to evaluate the PDF.
+
+        Returns
+        -------
+        DType | NDArray[DType]
+            The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
-        X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
 
-        return np.where(
-            self.scale <= X, (self.shape * (self.scale**self.shape)) / X ** (self.shape + dtype(1)), dtype(0.0)
-        )
+        X = np.asarray(X, dtype=self.dtype)
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -101,14 +109,20 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where((P >= 0) & (P <= 1), self.scale * (dtype(1) - P) ** (dtype(-1.0) / self.shape), dtype(np.nan))
+        result = np.where((P >= 0) & (P <= 1), self.scale * (dtype(1) - P) ** (dtype(-1.0) / self.shape), dtype(np.nan))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
@@ -129,19 +143,25 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(
+        result = np.where(
             self.scale <= X,
             np.log(self.shape) + self.shape * np.log(self.scale) - (dtype(1) + self.shape) * np.log(X),
             dtype(-np.inf),
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def _dlog_shape(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`shape` parameter.
 
@@ -162,14 +182,20 @@ def _dlog_shape(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`shape` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(self.scale <= X, dtype(1.0) / self.shape + np.log(self.scale) - np.log(X), dtype(0.0))
+        result = np.where(self.scale <= X, dtype(1.0) / self.shape + np.log(self.scale) - np.log(X), dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_scale(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`scale` parameter.
@@ -190,14 +216,19 @@ def _dlog_scale(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The gradient of the lpdf with respect to :attr:`scale` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(self.scale <= X, self.shape / self.scale, dtype(0.0))
+        result = np.where(self.scale <= X, self.shape / self.scale, dtype(0.0))
+        if is_scalar:
+            return result[()]
+        return result
 
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
@@ -216,8 +247,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -232,6 +265,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/uniform.py b/rework_pysatl_mpest/distributions/uniform.py
index 03cf75c..b2ccd1d 100644
--- a/rework_pysatl_mpest/distributions/uniform.py
+++ b/rework_pysatl_mpest/distributions/uniform.py
@@ -97,17 +97,13 @@ def pdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
-        X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
 
-        return np.where(
-            (self.left_border <= X) & (self.right_border >= X),
-            dtype(1.0) / (self.right_border - self.left_border),
-            dtype(0.0),
-        )
+        X = np.asarray(X, dtype=self.dtype)
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -128,16 +124,23 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
-        return np.where(
+        result = np.where(
             (P >= 0) & (P <= 1), self.left_border + P * (self.right_border - self.left_border), dtype(np.nan)
         )
 
+        if is_scalar:
+            return result[()]
+        return result
+
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
 
@@ -158,15 +161,22 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_range = (self.left_border <= X) & (self.right_border >= X)
         valid_dist = self.right_border > self.left_border
-        return np.where(in_range & valid_dist, -np.log(self.right_border - self.left_border), dtype(-np.inf))
+        result = np.where(in_range & valid_dist, -np.log(self.right_border - self.left_border), dtype(-np.inf))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_left_border(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`left_border` parameter.
@@ -177,13 +187,29 @@ def _dlog_left_border(self, X):
 
         where :math:`\\alpha` is the left_border parameter and :math:`\\beta` is the
         right_border parameter. The derivative is non-zero only for `left_border <= X <= right_border`.
+
+        Parameters
+        ----------
+        X : ArrayLike
+            The input data points.
+
+        Returns
+        -------
+        DType | NDArray[DType]
+            The gradient of the lpdf with respect to :attr:`left_border` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_range = (self.left_border <= X) & (self.right_border >= X)
-        return np.where(in_range, dtype(1.0) / (self.right_border - self.left_border), dtype(0.0))
+        result = np.where(in_range, dtype(1.0) / (self.right_border - self.left_border), dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_right_border(self, X):
         """Partial derivative of the lpdf w.r.t. the :attr:`right_border` parameter.
@@ -194,13 +220,29 @@ def _dlog_right_border(self, X):
 
         where :math:`\\alpha` is the left_border parameter and :math:`\\beta` is the
         right_border parameter. The derivative is non-zero only for `left_border <= X <= right_border`.
+
+        Parameters
+        ----------
+        X : ArrayLike
+            The input data points.
+
+        Returns
+        -------
+        DType | NDArray[DType]
+            The gradient of the lpdf with respect to :attr:`right_border` for each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         in_range = (self.left_border <= X) & (self.right_border >= X)
-        return np.where(in_range, dtype(-1.0) / (self.right_border - self.left_border), dtype(0.0))
+        result = np.where(in_range, dtype(-1.0) / (self.right_border - self.left_border), dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
@@ -219,8 +261,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -235,6 +279,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.asarray(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_pysatl_mpest/distributions/weibull.py b/rework_pysatl_mpest/distributions/weibull.py
index e06e449..7933164 100644
--- a/rework_pysatl_mpest/distributions/weibull.py
+++ b/rework_pysatl_mpest/distributions/weibull.py
@@ -91,21 +91,13 @@ def pdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
         X = np.asarray(X, dtype=self.dtype)
-        dtype = self.dtype
-
-        z = (X - self.loc) / self.scale
-
-        # PDF is 0 for x < loc, and handle cases where z=0 and shape<1
-        # which would lead to division by zero.
-        with np.errstate(divide="ignore", invalid="ignore"):
-            pdf_vals = (self.shape / self.scale) * np.power(z, self.shape - dtype(1)) * np.exp(-np.power(z, self.shape))
-
-        return np.where(self.loc <= X, np.nan_to_num(pdf_vals, nan=dtype(0.0), posinf=dtype(np.inf)), dtype(0.0))
+        return np.exp(self.lpdf(X))
 
     def ppf(self, P):
         """Percent Point Function (PPF) or quantile function.
@@ -124,15 +116,21 @@ def ppf(self, P):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The PPF values corresponding to each probability in :attr:`P`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(P)
         P = np.asarray(P, dtype=self.dtype)
         dtype = self.dtype
 
         ppf_vals = self.loc + self.scale * np.power(-np.log(dtype(1) - P), dtype(1.0) / self.shape)
-        return np.where((P >= 0) & (P <= 1), ppf_vals, dtype(np.nan))
+        result = np.where((P >= 0) & (P <= 1), ppf_vals, dtype(np.nan))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def lpdf(self, X):
         """Log of the Probability Density Function (LPDF).
@@ -152,34 +150,55 @@ def lpdf(self, X):
 
         Returns
         -------
-        NDArray[DType]
+        DType | NDArray[DType]
             The log-PDF values corresponding to each point in :attr:`X`.
+            Return a scalar when given a scalar, and to return an array when given an array.
         """
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         z = (X - self.loc) / self.scale
-        with np.errstate(divide="ignore"):
+        with np.errstate(divide="ignore", invalid="ignore"):
+            # Handle potential NaN from 0 * -inf when shape=1 and z=0
+            # This term's limit is 0, so we replace NaN with 0.
             lpdf_vals = (
-                np.log(self.shape) - np.log(self.scale) + (self.shape - dtype(1)) * np.log(z) - np.power(z, self.shape)
+                np.log(self.shape)
+                - np.log(self.scale)
+                + np.nan_to_num((self.shape - dtype(1)) * np.log(z), nan=dtype(0.0))
+                - np.power(z, self.shape)
             )
-        return np.where(self.loc < X, lpdf_vals, dtype(-np.inf))
+        result = np.where(self.loc < X, lpdf_vals, dtype(-np.inf))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_shape(self, X):
         """Partial derivative of the lpdf w.r.t. the shape parameter."""
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         z = (X - self.loc) / self.scale
         with np.errstate(divide="ignore", invalid="ignore"):
-            grad = dtype(1.0) / self.shape + np.log(z) - np.power(z, self.shape) * np.log(z)
-        return np.where(self.loc < X, np.nan_to_num(grad), dtype(0.0))
+            # Handle z^k * ln(z), which -> 0 as z -> 0.
+            # This prevents NaN from 0 * -inf.
+            grad = (
+                dtype(1.0) / self.shape + np.log(z) - np.nan_to_num(np.power(z, self.shape) * np.log(z), nan=dtype(0.0))
+            )
+        result = np.where(self.loc < X, np.nan_to_num(grad), dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_loc(self, X):
         """Partial derivative of the lpdf w.r.t. the loc parameter."""
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
@@ -188,17 +207,26 @@ def _dlog_loc(self, X):
             grad = -(self.shape - dtype(1)) / (X - self.loc) + (self.shape / self.scale) * np.power(
                 z, self.shape - dtype(1)
             )
-        return np.where(self.loc < X, np.nan_to_num(grad), dtype(0.0))
+        result = np.where(self.loc < X, np.nan_to_num(grad), dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def _dlog_scale(self, X):
         """Partial derivative of the lpdf w.r.t. the scale parameter."""
 
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
         dtype = self.dtype
 
         z = (X - self.loc) / self.scale
         grad = -self.shape / self.scale + (self.shape / self.scale) * np.power(z, self.shape)
-        return np.where(self.loc < X, grad, dtype(0.0))
+        result = np.where(self.loc < X, grad, dtype(0.0))
+
+        if is_scalar:
+            return result[()]
+        return result
 
     def log_gradients(self, X):
         """Calculates the gradients of the log-PDF w.r.t. its parameters.
@@ -215,7 +243,10 @@ def log_gradients(self, X):
             and each column corresponds to the gradient with respect to a
             specific optimizable parameter. The order of columns corresponds
             to the sorted order of :attr:`self.params_to_optimize`.
+            Returns a 1D array if X is a scalar.
         """
+
+        is_scalar = np.isscalar(X)
         X = np.asarray(X, dtype=self.dtype)
 
         gradient_calculators = {
@@ -231,6 +262,8 @@ def log_gradients(self, X):
 
         gradients = [gradient_calculators[param](X) for param in optimizable_params]
 
+        if is_scalar:
+            return np.array(gradients)
         return np.stack(gradients, axis=1)
 
     def generate(self, size: int):
diff --git a/rework_tests/unit/core/test_mixture.py b/rework_tests/unit/core/test_mixture.py
index 6a8d599..812483c 100644
--- a/rework_tests/unit/core/test_mixture.py
+++ b/rework_tests/unit/core/test_mixture.py
@@ -13,6 +13,7 @@
 import pytest
 from hypothesis import given
 from hypothesis import strategies as st
+from hypothesis.extra.numpy import arrays
 from rework_pysatl_mpest.core import MixtureModel
 from rework_pysatl_mpest.distributions import ContinuousDistribution, Exponential
 
@@ -291,45 +292,84 @@ def test_remove_component_with_invalid_index_raises_error(self, mixture_model: M
             mixture_model.remove_component(2)
 
 
+@pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestMixtureModelCalculations:
     """Tests for calculation methods like pdf, lpdf, etc."""
 
-    @pytest.mark.parametrize("X", [1.5, [1.5], np.array([1.0, 1.5, 6.0])])
-    def test_pdf_calculation(self, mixture_model: MixtureModel, X):
+    @given(x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
+    def test_pdf_calculation_for_array(self, x, dtype):
         """Tests the PDF calculation against the definition."""
 
-        dtype = mixture_model.dtype
+        mixture_model = MixtureModel(
+            components=[Exponential(loc=0.0, rate=1.0), Exponential(loc=5.0, rate=2.0)], weights=[0.5, 0.5], dtype=dtype
+        )
         c1, c2 = mixture_model.components
         w1, w2 = mixture_model.weights
 
-        expected_pdf = w1 * c1.pdf(X) + w2 * c2.pdf(X)
-        calculated_pdf = mixture_model.pdf(X)
+        expected_pdf = w1 * c1.pdf(x) + w2 * c2.pdf(x)
+        calculated_pdf = mixture_model.pdf(x)
 
         assert calculated_pdf.dtype == dtype
-        if not np.isscalar(X):
+        if not np.isscalar(x):
             assert isinstance(calculated_pdf, np.ndarray)
 
-        np.testing.assert_allclose(calculated_pdf, expected_pdf, rtol=np.finfo(dtype).eps)
+        np.testing.assert_allclose(calculated_pdf, expected_pdf, atol=np.finfo(dtype).eps)
 
-    @pytest.mark.parametrize("X", [1.5, [1.5], np.array([1.0, 1.5, 6.0])])
-    def test_lpdf_calculation(self, mixture_model: MixtureModel, X):
-        """Tests the LPDF calculation against the definition."""
+    @given(x=st.floats(-1e6, 1e6))
+    def test_pdf_calculation_for_scalar(self, x, dtype):
+        """Tests the PDF calculation against the definition."""
 
-        dtype = mixture_model.dtype
+        mixture_model = MixtureModel(
+            components=[Exponential(loc=0.0, rate=1.0), Exponential(loc=5.0, rate=2.0)], weights=[0.5, 0.5], dtype=dtype
+        )
         c1, c2 = mixture_model.components
         w1, w2 = mixture_model.weights
 
-        expected_lpdf = np.log(w1 * c1.pdf(X) + w2 * c2.pdf(X))
-        calculated_lpdf = mixture_model.lpdf(X)
+        expected_pdf = w1 * c1.pdf(x) + w2 * c2.pdf(x)
+        calculated_pdf = mixture_model.pdf(x)
+
+        assert np.isscalar(calculated_pdf)
+        assert isinstance(calculated_pdf, dtype)
+        np.testing.assert_allclose(calculated_pdf, expected_pdf, atol=np.finfo(dtype).eps)
+
+    @given(x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
+    def test_lpdf_calculation_for_array(self, x, dtype):
+        """Tests the LPDF calculation against the definition."""
+
+        mixture_model = MixtureModel(
+            components=[Exponential(loc=0.0, rate=1.0), Exponential(loc=5.0, rate=2.0)], weights=[0.5, 0.5], dtype=dtype
+        )
+
+        expected_pdf = mixture_model.pdf(x)
+        calculated_lpdf = mixture_model.lpdf(x)
+
+        if not np.isscalar(x):
+            assert isinstance(calculated_lpdf, np.ndarray)
 
-        assert isinstance(calculated_lpdf, np.ndarray)
         assert calculated_lpdf.dtype == dtype
-        np.testing.assert_allclose(calculated_lpdf, expected_lpdf, rtol=np.finfo(dtype).eps)
+        np.testing.assert_allclose(np.exp(calculated_lpdf), expected_pdf, atol=np.finfo(dtype).eps)
+
+    @given(x=st.floats(-1e6, 1e6))
+    def test_lpdf_calculation_for_scalar(self, x, dtype):
+        """Tests the LPDF calculation against the definition."""
+
+        mixture_model = MixtureModel(
+            components=[Exponential(loc=0.0, rate=1.0), Exponential(loc=5.0, rate=2.0)], weights=[0.5, 0.5], dtype=dtype
+        )
 
-    def test_loglikelihood_calculation(self, mixture_model: MixtureModel):
+        expected_pdf = mixture_model.pdf(x)
+        calculated_lpdf = mixture_model.lpdf(x)
+
+        assert np.isscalar(calculated_lpdf)
+        assert calculated_lpdf.dtype == dtype
+        np.testing.assert_allclose(np.exp(calculated_lpdf), expected_pdf, atol=np.finfo(dtype).eps)
+
+    def test_loglikelihood_calculation(self, dtype):
         """Tests that loglikelihood is the sum of LPDF values."""
 
-        dtype = mixture_model.dtype
+        mixture_model = MixtureModel(
+            components=[Exponential(loc=0.0, rate=1.0), Exponential(loc=5.0, rate=2.0)], weights=[0.5, 0.5], dtype=dtype
+        )
         X = np.array([1.0, 1.5, 6.0])
         expected_loglikelihood = np.sum(mixture_model.lpdf(X))
         calculated_loglikelihood = mixture_model.loglikelihood(X)
diff --git a/rework_tests/unit/distributions/test_beta.py b/rework_tests/unit/distributions/test_beta.py
index 9887daf..4a5782a 100644
--- a/rework_tests/unit/distributions/test_beta.py
+++ b/rework_tests/unit/distributions/test_beta.py
@@ -50,8 +50,8 @@ def load_r_test_cases():
 
 
 @st.composite
-def st_params_and_x_for_grad(draw):
-    """Generator of valid params with x for gradient test"""
+def st_params_and_array_x_for_grad(draw):
+    """Generator of valid params with an array x for gradient test"""
     shape1 = draw(st.floats(0.1, 100))
     shape2 = draw(st.floats(0.1, 100))
     left = draw(st.floats(-100, 100))
@@ -61,7 +61,7 @@ def st_params_and_x_for_grad(draw):
     x = draw(
         arrays(
             np.float64,
-            shape=draw(st.integers(1, 5)),
+            shape=draw(st.integers(2, 5)),
             elements=st.floats(
                 min_value=left + margin, max_value=right - margin, allow_nan=False, allow_infinity=False
             ),
@@ -70,6 +70,19 @@ def st_params_and_x_for_grad(draw):
     return (shape1, shape2, left, right, x)
 
 
+@st.composite
+def st_params_and_scalar_x_for_grad(draw):
+    """Generator of valid params with a scalar x for gradient test"""
+    shape1 = draw(st.floats(0.1, 100))
+    shape2 = draw(st.floats(0.1, 100))
+    lower = draw(st.floats(-100, 100))
+    upper = draw(st.floats(lower + 0.1, lower + 100))
+
+    margin = 1e-2
+    x = draw(st.floats(min_value=lower + margin, max_value=upper - margin, allow_nan=False, allow_infinity=False))
+    return (shape1, shape2, lower, upper, x)
+
+
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestBetaInitialization:
     """Tests for the __init__ method and basic properties."""
@@ -161,8 +174,8 @@ class TestBetaPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_pdf_properties(self, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
         alpha, beta, left_border, right_border = 1.0, 1.0, 2.9, 10.0
         dist = Beta(alpha, beta, left_border, right_border, dtype=dtype)
@@ -172,6 +185,18 @@ def test_pdf_properties(self, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type and shape."""
+
+        alpha, beta, left_border, right_border = 1.0, 1.0, 2.9, 10.0
+        dist = Beta(alpha, beta, left_border, right_border, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @pytest.mark.parametrize("x,shape1,shape2,left_border,right_border,expected_pdf", load_r_test_cases())
     def test_pdf_against_R(self, x, shape1, shape2, left_border, right_border, expected_pdf):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -203,8 +228,8 @@ class TestBetaLPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(params=st_valid_params(), x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_lpdf_return_type_and_shape(self, params, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, params, x, dtype):
+        """Tests that for an array input, the LPDF returns an array with the correct type and shape."""
 
         shape1, shape2, left_border, right_border = params
         dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
@@ -213,6 +238,17 @@ def test_lpdf_return_type_and_shape(self, params, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(params=st_valid_params(), x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, params, x, dtype):
+        """Tests that for a scalar input, the LPDF returns a scalar with the correct type and shape."""
+
+        alpha, beta, left_border, right_border = params
+        dist = Beta(alpha, beta, left_border, right_border, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(params=st_valid_params(), x=st.floats(1e-6, 1e6))
     def test_lpdf_against_scipy(self, params, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
@@ -244,8 +280,8 @@ class TestBetaPPF:
     @given(
         params=st_valid_params(), p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True))
     )
-    def test_ppf_return_type_and_shape(self, params, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, params, p, dtype):
+        """Tests that for an array input, the PPDF returns an array with the correct type and shape."""
 
         shape1, shape2, left_border, right_border = params
         dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
@@ -254,6 +290,17 @@ def test_ppf_return_type_and_shape(self, params, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(params=st_valid_params(), p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, params, p, dtype):
+        """Tests that for a scalar input, the PPDF returns a scalar with the correct type and shape."""
+
+        alpha, beta, left_border, right_border = params
+        dist = Beta(alpha, beta, left_border, right_border, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(params=st_valid_params(), p=st.floats(0, 1))
     def test_ppf_against_scipy(self, params, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -278,8 +325,8 @@ class TestBetaGradients:
 
     h = 1e-6
 
-    @given(params_x=st_params_and_x_for_grad())
-    def test_dlog_shape1_numerical(self, params_x, dtype):
+    @given(params_x=st_params_and_array_x_for_grad())
+    def test_dlog_shape1_numerical_for_array_input(self, params_x, dtype):
         """Checks the analytical gradient for 'shape1' against a numerical approximation."""
 
         shape1, shape2, left_border, right_border, x = params_x
@@ -298,9 +345,22 @@ def test_dlog_shape1_numerical(self, params_x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
-    @given(params_x=st_params_and_x_for_grad())
-    def test_dlog_shape2_numerical(self, params_x, dtype):
+    @given(params_x=st_params_and_scalar_x_for_grad())
+    def test_dlog_shape1_for_scalar_input(self, params_x, dtype):
+        """Checks that the gradient for 'shape1' for a scalar input returns scalar."""
+
+        shape1, shape2, left_border, right_border, x = params_x
+
+        dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
+        analytical_grad = dist._dlog_alpha(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
+    @given(params_x=st_params_and_array_x_for_grad())
+    def test_dlog_shape2_numerical_for_array_input(self, params_x, dtype):
         """Checks the analytical gradient for 'shape2' against a numerical approximation."""
+
         shape1, shape2, left_border, right_border, x = params_x
 
         dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
@@ -317,9 +377,22 @@ def test_dlog_shape2_numerical(self, params_x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
-    @given(params_x=st_params_and_x_for_grad())
-    def test_dlog_left_border_numerical(self, params_x, dtype):
+    @given(params_x=st_params_and_scalar_x_for_grad())
+    def test_dlog_shape2_for_scalar_input(self, params_x, dtype):
+        """Checks that the gradient for 'shape2' for a scalar input returns scalar."""
+
+        shape1, shape2, left_border, right_border, x = params_x
+
+        dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
+        analytical_grad = dist._dlog_beta(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
+    @given(params_x=st_params_and_array_x_for_grad())
+    def test_dlog_left_border_numerical_for_array_input(self, params_x, dtype):
         """Checks the analytical gradient for 'left_border' against a numerical approximation."""
+
         shape1, shape2, left_border, right_border, x = params_x
 
         dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
@@ -336,8 +409,20 @@ def test_dlog_left_border_numerical(self, params_x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
-    @given(params_x=st_params_and_x_for_grad())
-    def test_dlog_right_border_numerical(self, params_x, dtype):
+    @given(params_x=st_params_and_scalar_x_for_grad())
+    def test_dlog_left_border_for_scalar_input(self, params_x, dtype):
+        """Checks that the gradient for 'left_border' for a scalar input returns scalar."""
+
+        shape1, shape2, left_border, right_border, x = params_x
+
+        dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
+        analytical_grad = dist._dlog_left_border(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
+    @given(params_x=st_params_and_array_x_for_grad())
+    def test_dlog_right_border_numerical_for_array_input(self, params_x, dtype):
         """Checks the analytical gradient for 'right_border' against a numerical approximation."""
         shape1, shape2, left_border, right_border, x = params_x
 
@@ -355,6 +440,18 @@ def test_dlog_right_border_numerical(self, params_x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(params_x=st_params_and_scalar_x_for_grad())
+    def test_dlog_right_border_for_scalar_input(self, params_x, dtype):
+        """Checks that the gradient for 'right_border' for a scalar input returns scalar."""
+
+        shape1, shape2, left_border, right_border, x = params_x
+
+        dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
+        analytical_grad = dist._dlog_right_border(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_shape_col, expected_params",
         [
@@ -392,6 +489,19 @@ def test_log_gradients_structure(self, fixed_params, expected_shape_col, expecte
             idx = sorted(expected_params).index("right_border")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_right_border(x))
 
+    @given(params_x=st_params_and_scalar_x_for_grad())
+    def test_log_gradients_for_scalar_input(self, params_x, dtype):
+        """Checks that the log_gradients for a scalar input returns 1D-array."""
+
+        shape1, shape2, left_border, right_border, x = params_x
+        dist = Beta(shape1, shape2, left_border, right_border, dtype=dtype)
+
+        gradients = dist.log_gradients(x)
+
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestBetaGenerate:
diff --git a/rework_tests/unit/distributions/test_cauchy.py b/rework_tests/unit/distributions/test_cauchy.py
index 098d591..31badfe 100644
--- a/rework_tests/unit/distributions/test_cauchy.py
+++ b/rework_tests/unit/distributions/test_cauchy.py
@@ -81,10 +81,9 @@ class TestCauchyPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_pdf_properties(self, loc, scale, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, loc, scale, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
-        loc, scale = 0.0, 1.0
         dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
         pdf_values = dist.pdf(x)
         assert isinstance(pdf_values, np.ndarray)
@@ -92,6 +91,17 @@ def test_pdf_properties(self, loc, scale, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, loc, scale, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        dist = Cauchy(loc, scale, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @given(loc=st_loc, scale=st_scale, x=st.floats(1e-6, 1e6))
     def test_pdf_against_scipy(self, loc, scale, x):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -111,13 +121,13 @@ def test_pdf_integral_is_one(self, loc, scale):
         np.testing.assert_allclose(1.0, integral, rtol=1e-5)
 
 
-class TestLogCauchyPDF:
+class TestCauchyLPDF:
     """Tests for the lpdf (log-PDF) method using hypothesis."""
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_lpdf_return_type_and_shape(self, loc, scale, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
 
         dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -125,6 +135,16 @@ def test_lpdf_return_type_and_shape(self, loc, scale, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(loc=st_loc, scale=st_scale, x=st.floats(1e-6, 1e6))
     def test_lpdf_against_scipy(self, loc, scale, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
@@ -142,8 +162,8 @@ class TestCauchyPPF:
     @given(
         loc=st_loc, scale=st_scale, p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True))
     )
-    def test_ppf_return_type_and_shape(self, loc, scale, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
 
         dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -151,6 +171,16 @@ def test_ppf_return_type_and_shape(self, loc, scale, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(loc=st_loc, scale=st_scale, p=st.floats(0, 1, exclude_max=True, exclude_min=True))
     def test_ppf_against_scipy(self, loc, scale, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -175,8 +205,8 @@ class TestCauchyGradients:
     h = 1e-6
 
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_loc_numerical(self, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'loc' against a numerical approximation."""
+    def test_dlog_loc_numerical_for_array_input(self, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'loc' against a numerical approximation for array input."""
 
         assume(np.all(x > (loc + self.h)))
 
@@ -194,9 +224,21 @@ def test_dlog_loc_numerical(self, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_loc_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the gradient for 'loc' for a scalar input returns a scalar."""
+
+        assume(x > loc + self.h)
+
+        dist = Cauchy(loc, scale, dtype=dtype)
+        analytical_grad = dist._dlog_loc(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_scale_numerical(self, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'scale' against a numerical approximation."""
+    def test_dlog_scale_numerical_for_array_input(self, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'scale' against a numerical approximation for array input."""
 
         assume(np.all(x > (loc + self.h)))
 
@@ -214,6 +256,18 @@ def test_dlog_scale_numerical(self, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_scale_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the gradient for 'scale' for a scalar input returns a scalar."""
+
+        assume(x > loc + self.h)
+
+        dist = Cauchy(loc, scale, dtype=dtype)
+        analytical_grad = dist._dlog_scale(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_shape_col, expected_params",
         [([], 2, ["loc", "scale"]), (["loc"], 1, ["scale"]), (["scale"], 1, ["loc"]), (["loc", "scale"], 0, [])],
@@ -239,6 +293,17 @@ def test_log_gradients_structure(self, fixed_params, expected_shape_col, expecte
             idx = sorted(expected_params).index("scale")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_scale(x))
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_log_gradients_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        dist = Cauchy(loc, scale, dtype=dtype)
+        gradients = dist.log_gradients(x)
+
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestCauchyGenerate:
@@ -271,6 +336,23 @@ def test_generate_negative_size(self, size, dtype):
         with pytest.raises(ValueError):
             dist.generate(size=size)
 
+    def test_generate_statistical_properties(self, dtype):
+        """Tests if the generated samples have the correct statistical properties (median)."""
+
+        np.random.seed(123)
+        random.seed(123)
+        loc, scale = 5.0, 2.0
+        dist = Cauchy(loc=loc, scale=scale, dtype=dtype)
+        size = 5000
+
+        samples = dist.generate(size=size)
+
+        # The mean and variance of the Cauchy distribution are undefined.
+        # The median is equal to the location parameter 'loc'.
+        theoretical_median = loc
+
+        assert np.median(samples) == pytest.approx(theoretical_median, rel=0.1)
+
     def test_generate_kolmogorov_smirnov(self, dtype):
         """Performs a Kolmogorov-Smirnov test to check if samples fit the distribution."""
 
diff --git a/rework_tests/unit/distributions/test_exponential.py b/rework_tests/unit/distributions/test_exponential.py
index c1a05a3..2f4d774 100644
--- a/rework_tests/unit/distributions/test_exponential.py
+++ b/rework_tests/unit/distributions/test_exponential.py
@@ -81,8 +81,8 @@ class TestExponentialPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, rate=st_rate, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_pdf_properties(self, loc, rate, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, loc, rate, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
         dist = Exponential(loc=loc, rate=rate, dtype=dtype)
         pdf_values = dist.pdf(x)
@@ -91,6 +91,17 @@ def test_pdf_properties(self, loc, rate, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, rate=st_rate, x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, loc, rate, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        dist = Exponential(loc, rate, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @given(loc=st_loc, rate=st_rate, x=st.floats(1e-6, 1e6))
     def test_pdf_against_scipy(self, loc, rate, x):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -124,8 +135,8 @@ class TestExponentialLPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, rate=st_rate, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_lpdf_return_type_and_shape(self, loc, rate, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, loc, rate, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
 
         dist = Exponential(loc=loc, rate=rate, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -133,6 +144,16 @@ def test_lpdf_return_type_and_shape(self, loc, rate, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, rate=st_rate, x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, loc, rate, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        dist = Exponential(loc, rate, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(loc=st_loc, rate=st_rate, x=st.floats(1e-6, 1e6))
     def test_lpdf_against_scipy(self, loc, rate, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
@@ -160,8 +181,8 @@ class TestExponentialPPF:
     @given(
         loc=st_loc, rate=st_rate, p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True))
     )
-    def test_ppf_return_type_and_shape(self, loc, rate, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, loc, rate, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
 
         dist = Exponential(loc=loc, rate=rate, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -169,6 +190,16 @@ def test_ppf_return_type_and_shape(self, loc, rate, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, rate=st_rate, p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, loc, rate, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        dist = Exponential(loc=loc, rate=rate, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(loc=st_loc, rate=st_rate, p=st.floats(0, 1, exclude_max=True, exclude_min=True))
     def test_ppf_against_scipy(self, loc, rate, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -193,8 +224,8 @@ class TestExponentialGradients:
     h = 1e-6
 
     @given(loc=st_loc, rate=st_rate, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_loc_numerical(self, loc, rate, x, dtype):
-        """Checks the analytical gradient for 'loc' against a numerical approximation."""
+    def test_dlog_loc_numerical_for_array_input(self, loc, rate, x, dtype):
+        """Checks the analytical gradient for 'loc' against a numerical approximation for array input."""
 
         assume(np.all(x > (loc + self.h)))
 
@@ -212,9 +243,20 @@ def test_dlog_loc_numerical(self, loc, rate, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(loc=st_loc, rate=st_rate, x=st.floats(1e-3, 1e3))
+    def test_dlog_loc_for_scalar_input(self, loc, rate, x, dtype):
+        """Checks that the gradient for 'loc' for a scalar input returns a scalar."""
+
+        assume(x > (loc + self.h))
+
+        dist = Exponential(loc, rate, dtype=dtype)
+        analytical_grad = dist._dlog_loc(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @given(loc=st_loc, rate=st_rate, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_rate_numerical(self, loc, rate, x, dtype):
-        """Checks the analytical gradient for 'rate' against a numerical approximation."""
+    def test_dlog_rate_numerical_for_array_input(self, loc, rate, x, dtype):
+        """Checks the analytical gradient for 'rate' against a numerical approximation for array input."""
 
         assume(np.all(x > (loc + self.h)))
 
@@ -232,6 +274,17 @@ def test_dlog_rate_numerical(self, loc, rate, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(loc=st_loc, rate=st_rate, x=st.floats(1e-3, 1e3))
+    def test_dlog_rate_for_scalar_input(self, loc, rate, x, dtype):
+        """Checks that the gradient for 'rate' for a scalar input returns a scalar."""
+
+        assume(x > (loc + self.h))
+
+        dist = Exponential(loc, rate, dtype=dtype)
+        analytical_grad = dist._dlog_rate(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_shape_col, expected_params",
         [([], 2, ["loc", "rate"]), (["loc"], 1, ["rate"]), (["rate"], 1, ["loc"]), (["loc", "rate"], 0, [])],
@@ -257,6 +310,16 @@ def test_log_gradients_structure(self, fixed_params, expected_shape_col, expecte
             idx = sorted(expected_params).index("rate")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_rate(x))
 
+    @given(loc=st_loc, rate=st_rate, x=st.floats(1e-3, 1e3))
+    def test_log_gradients_for_scalar_input(self, loc, rate, x, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        dist = Exponential(loc, rate, dtype=dtype)
+        gradients = dist.log_gradients(x)
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestExponentialGenerate:
diff --git a/rework_tests/unit/distributions/test_normal.py b/rework_tests/unit/distributions/test_normal.py
index 0ad3930..94b02e7 100644
--- a/rework_tests/unit/distributions/test_normal.py
+++ b/rework_tests/unit/distributions/test_normal.py
@@ -81,8 +81,8 @@ class TestNormalPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_pdf_properties(self, loc, scale, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, loc, scale, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
         dist = Normal(loc=loc, scale=scale, dtype=dtype)
         pdf_values = dist.pdf(x)
@@ -91,6 +91,17 @@ def test_pdf_properties(self, loc, scale, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, loc, scale, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        dist = Normal(loc, scale, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
     def test_pdf_against_scipy(self, loc, scale, x):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -114,8 +125,8 @@ class TestNormalLPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_lpdf_return_type_and_shape(self, loc, scale, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
 
         dist = Normal(loc=loc, scale=scale, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -123,6 +134,16 @@ def test_lpdf_return_type_and_shape(self, loc, scale, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        dist = Normal(loc=loc, scale=scale, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
     def test_lpdf_against_scipy(self, loc, scale, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
@@ -138,8 +159,8 @@ class TestNormalPPF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(loc=st_loc, scale=st_scale, p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1)))
-    def test_ppf_return_type_and_shape(self, loc, scale, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
 
         dist = Normal(loc=loc, scale=scale, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -147,6 +168,16 @@ def test_ppf_return_type_and_shape(self, loc, scale, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(loc=st_loc, scale=st_scale, p=st.floats(0, 1))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        dist = Normal(loc=loc, scale=scale, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(loc=st_loc, scale=st_scale, p=st.floats(1e-6, 1.0 - 1e-6))
     def test_ppf_against_scipy(self, loc, scale, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -156,6 +187,13 @@ def test_ppf_against_scipy(self, loc, scale, p):
         scipy_ppf = norm.ppf(p, loc=loc, scale=scale)
         np.testing.assert_allclose(custom_ppf, scipy_ppf, atol=1e-9)
 
+    @pytest.mark.parametrize("p_val", [-0.5, 1.1, 1.5])
+    def test_ppf_invalid_input(self, p_val):
+        """Tests that PPF returns NaN for probabilities outside the [0, 1] range."""
+
+        dist = Normal(loc=0.0, scale=1.0)
+        assert np.isnan(dist.ppf(p_val))
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestNormalGradients:
@@ -164,8 +202,8 @@ class TestNormalGradients:
     h = 1e-6
 
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(-1e3, 1e3)))
-    def test_dlog_loc_numerical(self, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'loc' against a numerical approximation."""
+    def test_dlog_loc_numerical_for_array_input(self, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'loc' against a numerical approximation for array input."""
 
         dist = Normal(loc, scale, dtype=dtype)
         analytical_grad = dist._dlog_loc(x)
@@ -181,9 +219,19 @@ def test_dlog_loc_numerical(self, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e3, 1e3))
+    def test_dlog_loc_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the gradient for 'loc' for a scalar input returns a scalar."""
+
+        dist = Normal(loc, scale, dtype=dtype)
+        analytical_grad = dist._dlog_loc(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @given(loc=st_loc, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(-1e3, 1e3)))
-    def test_dlog_scale_numerical(self, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'scale' against a numerical approximation."""
+    def test_dlog_scale_numerical_for_array_input(self, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'scale' against a numerical approximation for array input."""
 
         dist = Normal(loc, scale, dtype=dtype)
         analytical_grad = dist._dlog_scale(x)
@@ -199,11 +247,21 @@ def test_dlog_scale_numerical(self, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e3, 1e3))
+    def test_dlog_scale_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the gradient for 'scale' for a scalar input returns a scalar."""
+
+        dist = Normal(loc, scale, dtype=dtype)
+        analytical_grad = dist._dlog_scale(x)
+
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_cols, expected_params",
         [([], 2, ["loc", "scale"]), (["loc"], 1, ["scale"]), (["scale"], 1, ["loc"]), (["loc", "scale"], 0, [])],
     )
-    def test_log_gradients_structure(self, fixed_params, expected_cols, expected_params, dtype):
+    def test_log_gradients_structure_for_array_input(self, fixed_params, expected_cols, expected_params, dtype):
         """Tests the structure and content of log_gradients with various fixed parameters."""
 
         dist = Normal(loc=1.0, scale=2.0, dtype=dtype)
@@ -222,6 +280,17 @@ def test_log_gradients_structure(self, fixed_params, expected_cols, expected_par
             idx = sorted_params.index("scale")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_scale(x))
 
+    @given(loc=st_loc, scale=st_scale, x=st.floats(-1e3, 1e3))
+    def test_log_gradients_for_scalar_input(self, loc, scale, x, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        dist = Normal(loc, scale, dtype=dtype)
+        gradients = dist.log_gradients(x)
+
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestNormalGenerate:
diff --git a/rework_tests/unit/distributions/test_pareto.py b/rework_tests/unit/distributions/test_pareto.py
index 176c181..bb18a4d 100644
--- a/rework_tests/unit/distributions/test_pareto.py
+++ b/rework_tests/unit/distributions/test_pareto.py
@@ -116,8 +116,8 @@ class TestParetoPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e2, 1e2)))
-    def test_pdf_properties(self, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
         dist = Pareto(shape=1.0, scale=2.0, dtype=dtype)
         pdf_values = dist.pdf(x)
@@ -126,6 +126,18 @@ def test_pdf_properties(self, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(x=st.floats(-1e2, 1e2))
+    def test_pdf_properties_for_scalar_input(self, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        shape, scale = 1.0, 2.0
+        dist = Pareto(shape, scale, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @pytest.mark.parametrize("x,shape,scale,expected_pdf", load_r_test_cases())
     def test_pdf_against_R(self, shape, scale, x, expected_pdf):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -171,8 +183,8 @@ class TestParetoLPDF:
 
     @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
     @given(shape=st_shape, scale=st_scale, x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)))
-    def test_lpdf_return_type_and_shape(self, shape, scale, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, shape, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
 
         dist = Pareto(shape=shape, scale=scale, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -180,9 +192,20 @@ def test_lpdf_return_type_and_shape(self, shape, scale, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(shape=st_shape, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, shape, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        dist = Pareto(shape=shape, scale=scale, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(shape=st_shape, scale=st_scale, x=st.floats(1e-3, 1e3, allow_infinity=False, allow_nan=False))
     def test_lpdf_against_scipy(self, shape, scale, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
+
         assume(np.isfinite(pareto.logpdf(x, scale=scale, b=shape, loc=0.0)))
         dist = Pareto(shape=shape, scale=scale)
         custom_lpdf = dist.lpdf(x)
@@ -207,8 +230,8 @@ class TestParetoPPF:
         scale=st_scale,
         p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True)),
     )
-    def test_ppf_return_type_and_shape(self, shape, scale, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, shape, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
 
         dist = Pareto(shape=shape, scale=scale, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -216,6 +239,16 @@ def test_ppf_return_type_and_shape(self, shape, scale, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(shape=st_shape, scale=st_scale, p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, shape, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        dist = Pareto(shape=shape, scale=scale, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(shape=st_shape, scale=st_scale, p=st.floats(0, 1))
     def test_ppf_against_scipy(self, shape, scale, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -241,8 +274,8 @@ class TestParetoGradients:
 
     @settings(suppress_health_check=[HealthCheck.filter_too_much])
     @given(shape=st_shape, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_shape_numerical(self, shape, scale, x, dtype):
-        """Checks the analytical gradient for 'shape' against a numerical approximation."""
+    def test_dlog_shape_numerical_for_array_input(self, shape, scale, x, dtype):
+        """Checks the analytical gradient for 'shape' against a numerical approximation for array input."""
 
         assume(np.all(x > scale))
 
@@ -260,10 +293,20 @@ def test_dlog_shape_numerical(self, shape, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(shape=st_shape, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_shape_for_scalar_input(self, shape, scale, x, dtype):
+        """Checks that the gradient for 'shape' for a scalar input returns a scalar."""
+
+        assume(x > scale)
+        dist = Pareto(shape, scale, dtype=dtype)
+        analytical_grad = dist._dlog_shape(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @settings(suppress_health_check=[HealthCheck.filter_too_much])
     @given(shape=st_shape, scale=st_scale, x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)))
-    def test_dlog_scale_numerical(self, shape, scale, x, dtype):
-        """Checks the analytical gradient for 'scale' against a numerical approximation."""
+    def test_dlog_scale_numerical_for_array_input(self, shape, scale, x, dtype):
+        """Checks the analytical gradient for 'scale' against a numerical approximation for array input."""
 
         assume(np.all(x > scale + self.h))
 
@@ -281,6 +324,16 @@ def test_dlog_scale_numerical(self, shape, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(shape=st_shape, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_scale_for_scalar_input(self, shape, scale, x, dtype):
+        """Checks that the gradient for 'scale' for a scalar input returns a scalar."""
+
+        assume(x > scale + self.h)
+        dist = Pareto(shape, scale, dtype=dtype)
+        analytical_grad = dist._dlog_scale(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_shape_col, expected_params",
         [
@@ -311,6 +364,17 @@ def test_log_gradients_structure(self, fixed_params, expected_shape_col, expecte
             idx = sorted(expected_params).index("scale")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_scale(x))
 
+    @given(shape=st_shape, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_log_gradients_for_scalar_input(self, shape, scale, x, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        assume(x > scale + self.h)
+        dist = Pareto(shape, scale, dtype=dtype)
+        gradients = dist.log_gradients(x)
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestParetoGenerate:
diff --git a/rework_tests/unit/distributions/test_uniform.py b/rework_tests/unit/distributions/test_uniform.py
index 21a2406..ce18e78 100644
--- a/rework_tests/unit/distributions/test_uniform.py
+++ b/rework_tests/unit/distributions/test_uniform.py
@@ -21,6 +21,7 @@
 @st.composite
 def st_valid_border(draw):
     """Generates valid borders"""
+
     left_border = draw(st.floats(min_value=-1e3, max_value=1e3 - 1, allow_nan=False, allow_infinity=False))
     right_border = draw(
         st.floats(min_value=left_border + 1e-6, max_value=left_border + 1e3, allow_nan=False, allow_infinity=False)
@@ -56,6 +57,7 @@ def test_params_property(self, dtype):
 
     def test_invariant_violation(self, dtype):
         """Tests that initializing with a infinite borders or left border bigger right border  raises a ValueError."""
+
         with pytest.raises(ValueError, match="right_border parameter must be strictly greater than left_border"):
             Uniform(0.0, -1.0, dtype=dtype)
         with pytest.raises(ValueError, match="right_border parameter must be strictly greater than left_border"):
@@ -89,8 +91,9 @@ class TestUniformPDF:
         borders=st_valid_border(),
         x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)),
     )
-    def test_pdf_properties(self, borders, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, borders, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
+
         left_border, right_border = borders
 
         dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
@@ -100,9 +103,22 @@ def test_pdf_properties(self, borders, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        left_border, right_border = -1.0, 12.0
+        dist = Uniform(left_border, right_border, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @given(borders=st_valid_border(), x=st.floats(1e-6, 1e6))
     def test_pdf_against_scipy(self, borders, x):
         """Compares the custom PDF implementation against scipy's implementation."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border)
         custom_pdf = dist.pdf(x)
@@ -112,6 +128,7 @@ def test_pdf_against_scipy(self, borders, x):
     @given(borders=st_valid_border())
     def test_pdf_integral_is_one(self, borders):
         """Tests that the integral of the PDF over its support is equal to 1."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border)
         integral, error = quad(lambda x: dist.pdf(x).item(), left_border, right_border)
@@ -120,6 +137,7 @@ def test_pdf_integral_is_one(self, borders):
     @given(borders=st_valid_border(), x=st.floats(max_value=-1e9, allow_infinity=False))
     def test_pdf_outside_support(self, borders, x):
         """Tests that the PDF is zero for values not in range of parameters."""
+
         left_border, right_border = borders
         x_val = left_border - abs(x)
         dist = Uniform(left_border=left_border, right_border=right_border)
@@ -134,8 +152,9 @@ class TestUniformLPDF:
         borders=st_valid_border(),
         x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)),
     )
-    def test_lpdf_return_type_and_shape(self, borders, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, borders, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -143,9 +162,21 @@ def test_lpdf_return_type_and_shape(self, borders, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(borders=st_valid_border(), x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, borders, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        left_border, right_border = borders
+        dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(borders=st_valid_border(), x=st.floats(1e-6, 1e6))
     def test_lpdf_against_scipy(self, borders, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border)
         custom_lpdf = dist.lpdf(x)
@@ -155,6 +186,7 @@ def test_lpdf_against_scipy(self, borders, x):
     @given(borders=st_valid_border(), x=st.floats(min_value=1e-6))
     def test_lpdf_outside_support(self, borders, x):
         """Tests that the LPDF is -inf for values outside the support."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border)
         assert dist.lpdf(left_border - x) == -np.inf
@@ -169,8 +201,9 @@ class TestUniformPPF:
         borders=st_valid_border(),
         p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True)),
     )
-    def test_ppf_return_type_and_shape(self, borders, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, borders, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -178,9 +211,21 @@ def test_ppf_return_type_and_shape(self, borders, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(borders=st_valid_border(), p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, borders, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        left_border, right_border = borders
+        dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(borders=st_valid_border(), p=st.floats(0, 1))
     def test_ppf_against_scipy(self, borders, p):
         """Compares the custom PPF implementation against scipy's implementation."""
+
         left_border, right_border = borders
         dist = Uniform(left_border=left_border, right_border=right_border)
         custom_ppf = dist.ppf(p)
@@ -196,8 +241,9 @@ def test_ppf_invalid_input(self, p_val):
 
 
 @st.composite
-def st_valid_grad_input(draw):
-    """Generates valid borders to calculate gradient"""
+def st_valid_grad_input_array(draw):
+    """Generates valid borders to calculate gradient for an array of x."""
+
     left_border = draw(st.floats(min_value=-10.0, max_value=10.0, allow_nan=False, allow_infinity=False))
     right_border = draw(
         st.floats(min_value=left_border + 0.1, max_value=left_border + 20.0, allow_nan=False, allow_infinity=False)
@@ -217,15 +263,35 @@ def st_valid_grad_input(draw):
     return (left_border, right_border), x_values
 
 
+@st.composite
+def st_valid_grad_input_scalar(draw):
+    """Generates valid borders to calculate gradient for a scalar x."""
+
+    left_border = draw(st.floats(min_value=-10.0, max_value=10.0, allow_nan=False, allow_infinity=False))
+    right_border = draw(
+        st.floats(min_value=left_border + 0.1, max_value=left_border + 20.0, allow_nan=False, allow_infinity=False)
+    )
+
+    margin = 0.01
+    x_value = draw(
+        st.floats(
+            min_value=left_border + margin, max_value=right_border - margin, allow_nan=False, allow_infinity=False
+        )
+    )
+
+    return (left_border, right_border), x_value
+
+
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestUniformGradients:
     """Tests for gradient calculation methods."""
 
     h = 1e-6
 
-    @given(input_data=st_valid_grad_input())
-    def test_dlog_left_border_numerical(self, input_data, dtype):
-        """Checks the analytical gradient for 'left_border' against a numerical approximation."""
+    @given(input_data=st_valid_grad_input_array())
+    def test_dlog_left_border_numerical_for_array_input(self, input_data, dtype):
+        """Checks the analytical gradient for 'left_border' against a numerical approximation for array input."""
+
         borders, x = input_data
         left_border, right_border = borders
 
@@ -242,9 +308,22 @@ def test_dlog_left_border_numerical(self, input_data, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
-    @given(input_data=st_valid_grad_input())
-    def test_dlog_right_border_numerical(self, input_data, dtype):
-        """Checks the analytical gradient for 'right_border' against a numerical approximation."""
+    @given(input_data=st_valid_grad_input_scalar())
+    def test_dlog_left_border_for_scalar_input(self, input_data, dtype):
+        """Checks that the gradient for 'left_border' for a scalar input returns a scalar."""
+
+        borders, x = input_data
+        left_border, right_border = borders
+
+        dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
+        analytical_grad = dist._dlog_left_border(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
+    @given(input_data=st_valid_grad_input_array())
+    def test_dlog_right_border_numerical_for_array_input(self, input_data, dtype):
+        """Checks the analytical gradient for 'right_border' against a numerical approximation for array input."""
+
         borders, x = input_data
         left_border, right_border = borders
 
@@ -261,6 +340,18 @@ def test_dlog_right_border_numerical(self, input_data, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(input_data=st_valid_grad_input_scalar())
+    def test_dlog_right_border_for_scalar_input(self, input_data, dtype):
+        """Checks that the gradient for 'right_border' for a scalar input returns a scalar."""
+
+        borders, x = input_data
+        left_border, right_border = borders
+
+        dist = Uniform(left_border=left_border, right_border=right_border, dtype=dtype)
+        analytical_grad = dist._dlog_right_border(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_shape_col, expected_params",
         [
@@ -291,6 +382,18 @@ def test_log_gradients_structure(self, fixed_params, expected_shape_col, expecte
             idx = sorted(expected_params).index("right_border")
             np.testing.assert_allclose(gradients[:, idx], dist._dlog_right_border(x))
 
+    @given(input_data=st_valid_grad_input_scalar())
+    def test_log_gradients_for_scalar_input(self, input_data, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        borders, x = input_data
+        left_border, right_border = borders
+        dist = Uniform(left_border, right_border, dtype=dtype)
+        gradients = dist.log_gradients(x)
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestUniformGenerate:
diff --git a/rework_tests/unit/distributions/test_weibull.py b/rework_tests/unit/distributions/test_weibull.py
index b334fa6..c4f5e6a 100644
--- a/rework_tests/unit/distributions/test_weibull.py
+++ b/rework_tests/unit/distributions/test_weibull.py
@@ -86,10 +86,10 @@ class TestWeibullPDF:
         shape=st_shape,
         loc=st_loc,
         scale=st_scale,
-        x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)),
+        x=arrays(np.float64, st.integers(1, 10), elements=st.floats(-1e6, 1e6)),
     )
-    def test_pdf_properties(self, shape, loc, scale, x, dtype):
-        """Tests that the PDF is non-negative and has the correct return type and shape."""
+    def test_pdf_properties_for_array_input(self, shape, loc, scale, x, dtype):
+        """Tests that for an array input, the PDF returns a non-negative array with the correct type and shape."""
 
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
         pdf_values = dist.pdf(x)
@@ -98,6 +98,17 @@ def test_pdf_properties(self, shape, loc, scale, x, dtype):
         assert pdf_values.shape == x.shape
         assert np.all(pdf_values >= 0)
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_pdf_properties_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Tests that for a scalar input, the PDF returns a non-negative scalar with the correct type."""
+
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        pdf_value = dist.pdf(x)
+        assert np.isscalar(pdf_value)
+        assert isinstance(pdf_value, dtype)
+        assert pdf_value >= 0
+
     @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-6, 1e6))
     def test_pdf_against_scipy(self, shape, loc, scale, x):
         """Compares the custom PDF implementation against scipy's implementation."""
@@ -135,8 +146,8 @@ class TestWeibullLPDF:
         scale=st_scale,
         x=arrays(np.float64, st.integers(0, 10), elements=st.floats(-1e6, 1e6)),
     )
-    def test_lpdf_return_type_and_shape(self, shape, loc, scale, x, dtype):
-        """Tests the return type and shape of the lpdf method."""
+    def test_lpdf_return_type_and_shape_for_array_input(self, shape, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for array input."""
 
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
         lpdf_values = dist.lpdf(x)
@@ -144,6 +155,16 @@ def test_lpdf_return_type_and_shape(self, shape, loc, scale, x, dtype):
         assert lpdf_values.dtype == dtype
         assert lpdf_values.shape == x.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(-1e6, 1e6))
+    def test_lpdf_return_type_and_shape_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Tests the return type and shape of the lpdf method for scalar input."""
+
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        lpdf_value = dist.lpdf(x)
+        assert np.isscalar(lpdf_value)
+        assert isinstance(lpdf_value, dtype)
+
     @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-6, 1e6))
     def test_lpdf_against_scipy(self, shape, loc, scale, x):
         """Compares the custom LPDF implementation against scipy's implementation."""
@@ -173,8 +194,8 @@ class TestWeibullPPF:
         scale=st_scale,
         p=arrays(np.float64, st.integers(0, 10), elements=st.floats(0, 1, exclude_max=True)),
     )
-    def test_ppf_return_type_and_shape(self, shape, loc, scale, p, dtype):
-        """Tests the return type and shape of the ppf method."""
+    def test_ppf_return_type_and_shape_for_array_input(self, shape, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for array input."""
 
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
         ppf_values = dist.ppf(p)
@@ -182,6 +203,16 @@ def test_ppf_return_type_and_shape(self, shape, loc, scale, p, dtype):
         assert ppf_values.dtype == dtype
         assert ppf_values.shape == p.shape
 
+    @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, p=st.floats(0, 1, exclude_max=True))
+    def test_ppf_return_type_and_shape_for_scalar_input(self, shape, loc, scale, p, dtype):
+        """Tests the return type and shape of the ppf method for scalar input."""
+
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        ppf_value = dist.ppf(p)
+        assert np.isscalar(ppf_value)
+        assert isinstance(ppf_value, dtype)
+
     @given(shape=st_shape, loc=st_loc, scale=st_scale, p=st.floats(0.01, 1))
     def test_ppf_against_scipy(self, shape, loc, scale, p):
         """Compares the custom PPF implementation against scipy's implementation."""
@@ -211,8 +242,8 @@ class TestWeibullGradients:
         scale=st_scale,
         x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)),
     )
-    def test_dlog_shape_numerical(self, shape, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'shape' against a numerical approximation."""
+    def test_dlog_shape_numerical_for_array_input(self, shape, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'shape' against a numerical approximation for array input."""
 
         assume(np.all(x > loc))
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
@@ -227,14 +258,24 @@ def test_dlog_shape_numerical(self, shape, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_shape_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Checks that the gradient for 'shape' for a scalar input returns a scalar."""
+
+        assume(x > loc)
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        analytical_grad = dist._dlog_shape(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @given(
         shape=st_shape,
         loc=st_loc,
         scale=st_scale,
         x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)),
     )
-    def test_dlog_loc_numerical(self, shape, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'loc' against a numerical approximation."""
+    def test_dlog_loc_numerical_for_array_input(self, shape, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'loc' against a numerical approximation for array input."""
 
         assume(np.all(x > (loc + self.h)))
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
@@ -249,14 +290,24 @@ def test_dlog_loc_numerical(self, shape, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-3)
 
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_loc_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Checks that the gradient for 'loc' for a scalar input returns a scalar."""
+
+        assume(x > loc)
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        analytical_grad = dist._dlog_loc(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @given(
         shape=st_shape,
         loc=st_loc,
         scale=st_scale,
         x=arrays(np.float64, st.integers(1, 10), elements=st.floats(1e-3, 1e3)),
     )
-    def test_dlog_scale_numerical(self, shape, loc, scale, x, dtype):
-        """Checks the analytical gradient for 'scale' against a numerical approximation."""
+    def test_dlog_scale_numerical_for_array_input(self, shape, loc, scale, x, dtype):
+        """Checks the analytical gradient for 'scale' against a numerical approximation for array input."""
 
         assume(np.all(x > loc))
         dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
@@ -271,6 +322,16 @@ def test_dlog_scale_numerical(self, shape, loc, scale, x, dtype):
             numerical_grad = (lpdf_plus_h - lpdf_minus_h) / (2 * self.h)
             np.testing.assert_allclose(analytical_grad, numerical_grad, atol=1e-3, rtol=1e-3)
 
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_dlog_scale_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Checks that the gradient for 'scale' for a scalar input returns a scalar."""
+
+        assume(x > loc)
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        analytical_grad = dist._dlog_scale(x)
+        assert np.isscalar(analytical_grad)
+        assert isinstance(analytical_grad, dtype)
+
     @pytest.mark.parametrize(
         "fixed_params, expected_cols, expected_params",
         [
@@ -306,6 +367,18 @@ def test_log_gradients_structure(self, fixed_params, expected_cols, expected_par
         if "shape" in expected_params:
             np.testing.assert_allclose(gradients[:, sorted_params.index("shape")], dist._dlog_shape(x))
 
+    @given(shape=st_shape, loc=st_loc, scale=st_scale, x=st.floats(1e-3, 1e3))
+    def test_log_gradients_for_scalar_input(self, shape, loc, scale, x, dtype):
+        """Checks that the log_gradients for a scalar input returns a 1D-array."""
+
+        assume(x > loc)
+
+        dist = Weibull(shape=shape, loc=loc, scale=scale, dtype=dtype)
+        gradients = dist.log_gradients(x)
+        assert isinstance(gradients, np.ndarray)
+        assert gradients.dtype == dtype
+        assert gradients.ndim == 1
+
 
 @pytest.mark.parametrize("dtype", DTYPES_TO_TEST)
 class TestWeibullGenerate: