Fixed str parametrization for Negative Binomial distribution (#4187)

* Fixed repr parametrization for Negative Binomial distribution

* Added tests

Wrapped p and n with tt.as_tensor_variable

Author: Sayam753

pymc3/distributions/discrete.py

@@ -640,8 +640,11 @@ def __init__(self, mu=None, alpha=None, p=None, n=None, *args, **kwargs):
         self.mode = intX(tt.floor(mu))
 
     def get_mu_alpha(self, mu=None, alpha=None, p=None, n=None):
+        self._param_type = ["mu", "alpha"]
         if alpha is None:
             if n is not None:
+                self._param_type[1] = "n"
+                self.n = tt.as_tensor_variable(intX(n))
                 alpha = n
             else:
                 raise ValueError("Incompatible parametrization. Must specify either alpha or n.")
@@ -650,6 +653,8 @@ def get_mu_alpha(self, mu=None, alpha=None, p=None, n=None):
 
         if mu is None:
             if p is not None:
+                self._param_type[0] = "p"
+                self.p = tt.as_tensor_variable(floatX(p))
                 mu = alpha * (1 - p) / p
             else:
                 raise ValueError("Incompatible parametrization. Must specify either mu or p.")
@@ -720,6 +725,9 @@ def logp(self, value):
         # Return Poisson when alpha gets very large.
         return tt.switch(tt.gt(alpha, 1e10), Poisson.dist(self.mu).logp(value), negbinom)
 
+    def _distr_parameters_for_repr(self):
+        return self._param_type
+
 
 class Geometric(Discrete):
     R"""

pymc3/tests/test_distributions.py

@@ -1880,6 +1880,17 @@ def test_issue_3051(self, dims, dist_cls, kwargs):
         assert actual_a.shape == (X.shape[0],)
         pass
 
+    def test_issue_4186(self):
+        with pm.Model():
+            nb = pm.NegativeBinomial(
+                "nb", mu=pm.Normal("mu"), alpha=pm.Gamma("alpha", mu=6, sigma=1)
+            )
+        assert str(nb) == "nb ~ NegativeBinomial(mu=mu, alpha=alpha)"
+
+        with pm.Model():
+            nb = pm.NegativeBinomial("nb", p=pm.Uniform("p"), n=10)
+        assert str(nb) == "nb ~ NegativeBinomial(p=p, n=10)"
+
 
 def test_serialize_density_dist():
     def func(x):