Allow passing dims to Potential and Deterministic (#6576)

* Adding dims parameter in Potential

* Test case to check the dims in potential

Author: Raj-Parekh24

pymc/model.py

@@ -2041,7 +2041,7 @@ def Deterministic(name, var, model=None, dims=None):
     return var
 
 
-def Potential(name, var, model=None):
+def Potential(name, var, model=None, dims=None):
     """
     Add an arbitrary factor potential to the model likelihood
 
@@ -2135,7 +2135,7 @@ def Potential(name, var, model=None):
     model = modelcontext(model)
     var.name = model.name_for(name)
     model.potentials.append(var)
-    model.add_named_variable(var)
+    model.add_named_variable(var, dims)
 
     from pymc.printing import str_for_potential_or_deterministic
 

tests/test_model.py

@@ -1062,6 +1062,26 @@ def test_deterministic():
     assert model["y"] == y
 
 
+def test_determinsitic_with_dims():
+    """
+    Test to check the passing of dims to the potential
+    """
+    with pm.Model(coords={"observed": range(10)}) as model:
+        x = pm.Normal("x", 0, 1)
+        y = pm.Deterministic("y", x**2, dims=("observed",))
+    assert model.named_vars_to_dims == {"y": ("observed",)}
+
+
+def test_potential_with_dims():
+    """
+    Test to check the passing of dims to the potential
+    """
+    with pm.Model(coords={"observed": range(10)}) as model:
+        x = pm.Normal("x", 0, 1)
+        y = pm.Potential("y", x**2, dims=("observed",))
+    assert model.named_vars_to_dims == {"y": ("observed",)}
+
+
 def test_empty_model_representation():
     assert pm.Model().str_repr() == ""