From c1e17331458caec6b2cc3add066461338122e686 Mon Sep 17 00:00:00 2001
From: Juan Orduz <juan.orduz@wolt.com>
Date: Tue, 5 Nov 2024 21:08:59 +0100
Subject: [PATCH 1/2] suggest dimension improvement

---
 pymc/gp/hsgp_approx.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py
index cf434ebe3fe..4e0a403acce 100644
--- a/pymc/gp/hsgp_approx.py
+++ b/pymc/gp/hsgp_approx.py
@@ -428,8 +428,8 @@ def prior(
         self,
         name: str,
         X: TensorLike,
+        dims: str | None = None,
         hsgp_coeffs_dims: str | None = None,
-        gp_dims: str | None = None,
         *args,
         **kwargs,
     ):
@@ -444,10 +444,11 @@ def prior(
             Name of the random variable
         X: array-like
             Function input values.
+        dims: str, default None
+            Dimension name for the GP random variable.
         hsgp_coeffs_dims: str, default None
             Dimension name for the HSGP basis vectors.
-        gp_dims: str, default None
-            Dimension name for the GP random variable.
+
         """
         phi, sqrt_psd = self.prior_linearized(X)
         self._sqrt_psd = sqrt_psd
@@ -469,7 +470,7 @@ def prior(
             )
             f = self.mean_func(X) + phi @ self._beta
 
-        self.f = pm.Deterministic(name, f, dims=gp_dims)
+        self.f = pm.Deterministic(name, f, dims=dims)
         return self.f
 
     def _build_conditional(self, Xnew):
@@ -695,7 +696,9 @@ def prior_linearized(self, X: TensorLike):
         psd = self.scale * self.cov_func.power_spectral_density_approx(J)
         return (phi_cos, phi_sin), psd
 
-    def prior(self, name: str, X: TensorLike, dims: str | None = None):  # type: ignore[override]
+    def prior(
+        self, name: str, X: TensorLike, dims: str | None = None, hsgp_coeffs_dims: str | None = None
+    ):  # type: ignore[override]
         R"""
         Return the (approximate) GP prior distribution evaluated over the input locations `X`.
 
@@ -709,11 +712,13 @@ def prior(self, name: str, X: TensorLike, dims: str | None = None):  # type: ign
             Function input values.
         dims: None
             Dimension name for the GP random variable.
+        hsgp_coeffs_dims: str | None = None
+            Dimension name for the HSGPPeriodic basis vectors.
         """
         (phi_cos, phi_sin), psd = self.prior_linearized(X)
 
         m = self._m
-        self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m * 2 - 1))
+        self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m * 2 - 1), dims=hsgp_coeffs_dims)
         # The first eigenfunction for the sine component is zero
         # and so does not contribute to the approximation.
         f = (

From af142c4579be0e22e3b286ef88d224fc4485438b Mon Sep 17 00:00:00 2001
From: Juan Orduz <juan.orduz@wolt.com>
Date: Tue, 5 Nov 2024 21:25:47 +0100
Subject: [PATCH 2/2] make mypy happy

---
 pymc/gp/hsgp_approx.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py
index 4e0a403acce..e331aab8f5d 100644
--- a/pymc/gp/hsgp_approx.py
+++ b/pymc/gp/hsgp_approx.py
@@ -696,9 +696,9 @@ def prior_linearized(self, X: TensorLike):
         psd = self.scale * self.cov_func.power_spectral_density_approx(J)
         return (phi_cos, phi_sin), psd
 
-    def prior(
+    def prior(  # type: ignore[override]
         self, name: str, X: TensorLike, dims: str | None = None, hsgp_coeffs_dims: str | None = None
-    ):  # type: ignore[override]
+    ):
         R"""
         Return the (approximate) GP prior distribution evaluated over the input locations `X`.