diff --git a/numpyro/contrib/hsgp/approximation.py b/numpyro/contrib/hsgp/approximation.py
index e179790f5..155c1e25a 100644
--- a/numpyro/contrib/hsgp/approximation.py
+++ b/numpyro/contrib/hsgp/approximation.py
@@ -68,8 +68,8 @@ def linear_approximation(
 
 def hsgp_squared_exponential(
     x: ArrayLike,
-    alpha: float,
-    length: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     non_centered: bool = True,
@@ -90,8 +90,8 @@ def hsgp_squared_exponential(
            approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayLike x: input data
-    :param float alpha: amplitude of the squared exponential kernel
-    :param float length: length scale of the squared exponential kernel
+    :param ArrayLike alpha: amplitude of the squared exponential kernel
+    :param ArrayLike length: length scale of the squared exponential kernel
     :param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
         that the input data lies in the interval :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_E]`.
         We expect the approximation to be valid within this interval
@@ -117,8 +117,8 @@ def hsgp_squared_exponential(
 def hsgp_matern(
     x: ArrayLike,
     nu: float,
-    alpha: float,
-    length: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     non_centered: bool = True,
@@ -140,8 +140,8 @@ def hsgp_matern(
 
     :param ArrayLike x: input data
     :param float nu: smoothness parameter
-    :param float alpha: amplitude of the squared exponential kernel
-    :param float length: length scale of the squared exponential kernel
+    :param ArrayLike alpha: amplitude of the squared exponential kernel
+    :param ArrayLike length: length scale of the squared exponential kernel
     :param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
         that the input data lies in the interval :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.
         We expect the approximation to be valid within this interval
@@ -166,9 +166,9 @@ def hsgp_matern(
 
 def hsgp_rational_quadratic(
     x: ArrayLike,
-    alpha: float,
-    length: float,
-    scale_mixture: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
+    scale_mixture: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     non_centered: bool = True,
@@ -202,9 +202,9 @@ def hsgp_rational_quadratic(
            approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayLike x: input data
-    :param float alpha: amplitude of the Rational Quadratic kernel
-    :param float length: length scale of the Rational Quadratic kernel (scalar, isotropic only)
-    :param float scale_mixture: scale mixture parameter (α in the RQ kernel formula).
+    :param ArrayLike alpha: amplitude of the Rational Quadratic kernel
+    :param ArrayLike length: length scale of the Rational Quadratic kernel (scalar, isotropic only)
+    :param ArrayLike scale_mixture: scale mixture parameter (α in the RQ kernel formula).
         Controls the relative weighting of small-scale and large-scale variations.
         As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
     :param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
@@ -235,7 +235,7 @@ def hsgp_rational_quadratic(
 
 
 def hsgp_periodic_non_centered(
-    x: ArrayLike, alpha: float, length: float, w0: float, m: int
+    x: ArrayLike, alpha: ArrayLike, length: ArrayLike, w0: ArrayLike, m: int
 ) -> Array:
     """
     Low rank approximation for the periodic squared exponential kernel in the non-centered parametrization.
@@ -248,9 +248,9 @@ def hsgp_periodic_non_centered(
            approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayLike x: input data
-    :param float alpha: amplitude
-    :param float length: length scale
-    :param float w0: frequency of the periodic kernel
+    :param ArrayLike alpha: amplitude
+    :param ArrayLike length: length scale
+    :param ArrayLike w0: frequency of the periodic kernel
     :param int m: number of eigenvalues to compute and include in the approximation
     :return: the low-rank approximation linear model
     :rtype: Array
diff --git a/numpyro/contrib/hsgp/laplacian.py b/numpyro/contrib/hsgp/laplacian.py
index ae0cbbd6b..8e6b6da44 100644
--- a/numpyro/contrib/hsgp/laplacian.py
+++ b/numpyro/contrib/hsgp/laplacian.py
@@ -162,12 +162,12 @@ def eigenfunctions(x: ArrayLike, ell: float | list[float], m: int | list[int]) -
     )
 
 
-def eigenfunctions_periodic(x: ArrayLike, w0: float, m: int) -> tuple[Array, Array]:
+def eigenfunctions_periodic(x: ArrayLike, w0: ArrayLike, m: int) -> tuple[Array, Array]:
     """
     Basis functions for the approximation of the periodic kernel.
 
     :param ArrayLike x: The points at which to evaluate the eigenfunctions.
-    :param float w0: The frequency of the periodic kernel.
+    :param ArrayLike w0: The frequency of the periodic kernel.
     :param int m: The number of eigenfunctions to compute.
 
     .. note::
diff --git a/numpyro/contrib/hsgp/spectral_densities.py b/numpyro/contrib/hsgp/spectral_densities.py
index b9af76fdb..78cb45919 100644
--- a/numpyro/contrib/hsgp/spectral_densities.py
+++ b/numpyro/contrib/hsgp/spectral_densities.py
@@ -20,7 +20,7 @@ def align_param(dim, param):
 
 
 def spectral_density_squared_exponential(
-    dim: int, w: ArrayLike, alpha: float, length: float | ArrayLike
+    dim: int, w: ArrayLike, alpha: ArrayLike, length: ArrayLike
 ) -> Array:
     """
     Spectral density of the squared exponential kernel.
@@ -42,20 +42,20 @@ def spectral_density_squared_exponential(
 
     :param int dim: dimension
     :param ArrayLike w: frequency
-    :param float alpha: amplitude
-    :param float length: length scale
+    :param ArrayLike alpha: amplitude
+    :param ArrayLike length: length scale
     :return: spectral density value
     :rtype: Array
     """
     length = align_param(dim, length)
-    c = alpha * jnp.prod(jnp.sqrt(2 * jnp.pi) * length, axis=-1)
+    c = jnp.prod(jnp.sqrt(2 * jnp.pi) * length, axis=-1) * alpha
     e = jnp.exp(-0.5 * jnp.sum(w**2 * length**2, axis=-1))
     return c * e
 
 
 def spectral_density_matern(
-    dim: int, nu: float, w: ArrayLike, alpha: float, length: float | ArrayLike
-) -> float:
+    dim: int, nu: float, w: ArrayLike, alpha: ArrayLike, length: ArrayLike
+) -> Array:
     """
     Spectral density of the Matérn kernel.
 
@@ -78,10 +78,10 @@ def spectral_density_matern(
     :param int dim: dimension
     :param float nu: smoothness
     :param ArrayLike w: frequency
-    :param float alpha: amplitude
-    :param float length: length scale
+    :param ArrayLike alpha: amplitude
+    :param ArrayLike length: length scale
     :return: spectral density value
-    :rtype: float
+    :rtype: Array
     """  # noqa: E501
     length = align_param(dim, length)
     c1 = (
@@ -98,8 +98,8 @@ def spectral_density_matern(
 
 
 def diag_spectral_density_squared_exponential(
-    alpha: float,
-    length: float | list[float],
+    alpha: ArrayLike,
+    length: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     dim: int,
@@ -108,8 +108,8 @@ def diag_spectral_density_squared_exponential(
     Evaluates the spectral density of the squared exponential kernel at the first :math:`D \\times m^\\star`
     square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.
 
-    :param float alpha: amplitude of the squared exponential kernel
-    :param float length: length scale of the squared exponential kernel
+    :param ArrayLike alpha: amplitude of the squared exponential kernel
+    :param ArrayLike length: length scale of the squared exponential kernel
     :param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
         If a float or int, the same length is used in each dimension.
     :param int | list[int] m: The number of eigenvalues to compute for each dimension.
@@ -125,7 +125,7 @@ def _spectral_density(w):
             dim=dim,
             w=w,
             alpha=alpha,
-            length=length,  # ty: ignore[invalid-argument-type]
+            length=length,
         )
 
     sqrt_eigenvalues_ = sqrt_eigenvalues(ell=ell, m=m, dim=dim)  # dim x m
@@ -135,8 +135,8 @@ def _spectral_density(w):
 # TODO support length-D kernel hyperparameters
 def diag_spectral_density_matern(
     nu: float,
-    alpha: float,
-    length: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     dim: int,
@@ -146,8 +146,8 @@ def diag_spectral_density_matern(
     square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.
 
     :param float nu: smoothness parameter
-    :param float alpha: amplitude of the Matérn kernel
-    :param float length: length scale of the Matérn kernel
+    :param ArrayLike alpha: amplitude of the Matérn kernel
+    :param ArrayLike length: length scale of the Matérn kernel
     :param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
         If a float or int, the same length is used in each dimension.
     :param int | list[int] m: The number of eigenvalues to compute for each dimension.
@@ -191,9 +191,9 @@ def modified_bessel_second_kind(v, z):
 def spectral_density_rational_quadratic(
     dim: int,
     w: ArrayLike,
-    alpha: float,
-    length: float | ArrayLike,
-    scale_mixture: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
+    scale_mixture: ArrayLike,
 ) -> Array:
     """
     Spectral density of the Rational Quadratic kernel.
@@ -240,9 +240,9 @@ def spectral_density_rational_quadratic(
 
     :param int dim: dimension
     :param ArrayLike w: frequency
-    :param float alpha: amplitude (σ² in the spectral density formula)
-    :param float length: length scale (scalar, isotropic)
-    :param float scale_mixture: scale mixture parameter (α_mix in the RQ kernel formula).
+    :param ArrayLike alpha: amplitude (σ² in the spectral density formula)
+    :param ArrayLike length: length scale (scalar, isotropic)
+    :param ArrayLike scale_mixture: scale mixture parameter (α_mix in the RQ kernel formula).
         Controls the relative weighting of small-scale and large-scale variations.
         As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
     :return: spectral density value
@@ -304,9 +304,9 @@ def spectral_density_rational_quadratic(
 
 
 def diag_spectral_density_rational_quadratic(
-    alpha: float,
-    length: float | list[float],
-    scale_mixture: float,
+    alpha: ArrayLike,
+    length: ArrayLike,
+    scale_mixture: ArrayLike,
     ell: float | int | list[float | int],
     m: int | list[int],
     dim: int,
@@ -315,9 +315,9 @@ def diag_spectral_density_rational_quadratic(
     Evaluates the spectral density of the Rational Quadratic kernel at the first :math:`D \\times m^\\star`
     square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.
 
-    :param float alpha: amplitude of the Rational Quadratic kernel
-    :param float length: length scale of the Rational Quadratic kernel
-    :param float scale_mixture: scale mixture parameter (α in the RQ kernel formula).
+    :param ArrayLike alpha: amplitude of the Rational Quadratic kernel
+    :param ArrayLike length: length scale of the Rational Quadratic kernel
+    :param ArrayLike scale_mixture: scale mixture parameter (α in the RQ kernel formula).
         Controls the relative weighting of small-scale and large-scale variations.
         As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
     :param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
@@ -335,7 +335,7 @@ def _spectral_density(w):
             dim=dim,
             w=w,
             alpha=alpha,
-            length=length,  # ty: ignore[invalid-argument-type]
+            length=length,
             scale_mixture=scale_mixture,
         )
 
@@ -356,7 +356,9 @@ def modified_bessel_first_kind(v, z):
     return jnp.exp(jnp.abs(z)) * tfp.math.bessel_ive(v, z)
 
 
-def diag_spectral_density_periodic(alpha: float, length: float, m: int) -> Array:
+def diag_spectral_density_periodic(
+    alpha: ArrayLike, length: ArrayLike, m: int
+) -> Array:
     """
     Not actually a spectral density but these are used in the same
     way. These are simply the first `m` coefficients of the low rank
@@ -367,8 +369,8 @@ def diag_spectral_density_periodic(alpha: float, length: float, m: int) -> Array
         1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
            approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
-    :param float alpha: amplitude
-    :param float length: length scale
+    :param ArrayLike alpha: amplitude
+    :param ArrayLike length: length scale
     :param int m: number of eigenvalues
     :return: "spectral density" vector
     :rtype: Array