diff --git a/aemcmc/gibbs.py b/aemcmc/gibbs.py
index 4bab440..b7edb0b 100644
--- a/aemcmc/gibbs.py
+++ b/aemcmc/gibbs.py
@@ -1,10 +1,15 @@
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 
 import aesara
 import aesara.tensor as at
+from aesara.graph import optimize_graph
+from aesara.graph.unify import eval_if_etuple
 from aesara.ifelse import ifelse
+from aesara.tensor.math import Dot
 from aesara.tensor.random import RandomStream
-from aesara.tensor.var import TensorVariable, Variable
+from aesara.tensor.var import TensorVariable
+from etuples import etuple, etuplize
+from unification import unify, var
 
 from aemcmc.dists import (
     multivariate_normal_cong2017,
@@ -13,39 +18,128 @@
 )
 
 
-def update_beta_low_dimension(rng, omega, lambda2tau2_inv, X, z):
+def update_beta_low_dimension(
+    srng: RandomStream,
+    omega: TensorVariable,
+    lmbdatau_inv: TensorVariable,
+    X: TensorVariable,
+    z: TensorVariable,
+) -> TensorVariable:
     Q = X.T @ (omega[:, None] * X)
     indices = at.arange(Q.shape[1])
     Q = at.subtensor.set_subtensor(
         Q[indices, indices],
-        at.diag(Q) + lambda2tau2_inv,
+        at.diag(Q) + lmbdatau_inv,
     )
-    return multivariate_normal_rue2005(rng, X.T @ (omega * z), Q)
+    return multivariate_normal_rue2005(srng, X.T @ (omega * z), Q)
 
 
-def update_beta_high_dimension(rng, omega, lambda2tau2_inv, X, z):
-    return multivariate_normal_cong2017(rng, lambda2tau2_inv, omega, X, z)
+def update_beta_high_dimension(
+    srng: RandomStream,
+    omega: TensorVariable,
+    lmbdatau_inv: TensorVariable,
+    X: TensorVariable,
+    z: TensorVariable,
+) -> TensorVariable:
+    return multivariate_normal_cong2017(srng, lmbdatau_inv, omega, X, z)
 
 
-def update_beta(rng, omega, lambda2tau2_inv, X, z):
+def update_beta(
+    srng: RandomStream,
+    omega: TensorVariable,
+    lmbdatau_inv: TensorVariable,
+    X: TensorVariable,
+    z: TensorVariable,
+) -> TensorVariable:
     return ifelse(
         X.shape[1] > X.shape[0],
-        update_beta_high_dimension(rng, omega, lambda2tau2_inv, X, z),
-        update_beta_low_dimension(rng, omega, lambda2tau2_inv, X, z),
+        update_beta_high_dimension(srng, omega, lmbdatau_inv, X, z),
+        update_beta_low_dimension(srng, omega, lmbdatau_inv, X, z),
     )
 
 
+halfcauchy_1_lv, halfcauchy_2_lv = var(), var()
+zero_lv = var()
+horseshoe_pattern = etuple(
+    etuplize(at.random.normal),
+    var(),
+    var(),
+    var(),
+    zero_lv,
+    etuple(etuplize(at.mul), halfcauchy_1_lv, halfcauchy_2_lv),
+)
+
+
+def horseshoe_model(srng: TensorVariable) -> TensorVariable:
+    """Horseshoe shrinkage prior [1]_.
+
+    References
+    ----------
+    .. [1]: Carvalho, C. M., Polson, N. G., & Scott, J. G. (2010).
+            The horseshoe estimator for sparse signals.
+            Biometrika, 97(2), 465-480.
+
+    """
+    size = at.scalar("size", dtype="int32")
+    tau_rv = srng.halfcauchy(0, 1, size=1)
+    lmbda_rv = srng.halfcauchy(0, 1, size=size)
+    beta_rv = srng.normal(0, tau_rv * lmbda_rv, size=size)
+    return beta_rv
+
+
+def horseshoe_match(graph: TensorVariable) -> Tuple[TensorVariable, TensorVariable]:
+    graph_opt = optimize_graph(graph)
+    graph_et = etuplize(graph_opt)
+
+    s = unify(graph_et, horseshoe_pattern)
+    if s is False:
+        raise ValueError("Not a horseshoe prior.")
+
+    halfcauchy_1 = eval_if_etuple(s[halfcauchy_1_lv])
+    if halfcauchy_1.owner is None or not isinstance(
+        halfcauchy_1.owner.op, type(at.random.halfcauchy)
+    ):
+        raise ValueError(
+            "Not a horseshoe prior. One of the shrinkage parameters "
+            + "in your model is not half-Cauchy distributed."
+        )
+
+    halfcauchy_2 = eval_if_etuple(s[halfcauchy_2_lv])
+
+    if halfcauchy_2.owner is None or not isinstance(
+        halfcauchy_2.owner.op, type(at.random.halfcauchy)
+    ):
+        raise ValueError(
+            "Not a horseshoe prior. One of the shrinkage parameters "
+            + "in your model is not half-Cauchy distributed."
+        )
+
+    if halfcauchy_1.type.shape == (1,):
+        lmbda_rv = halfcauchy_2
+        tau_rv = halfcauchy_1
+    elif halfcauchy_2.type.shape == (1,):
+        lmbda_rv = halfcauchy_1
+        tau_rv = halfcauchy_2
+    else:
+        raise ValueError(
+            "Not a horseshoe prior. The global shrinkage parameter "
+            + "in your model must be one-dimensional."
+        )
+
+    return (lmbda_rv, tau_rv)
+
+
 def horseshoe_step(
     srng: RandomStream,
     beta: TensorVariable,
-    sigma2: TensorVariable,
-    lambda2_inv: TensorVariable,
-    tau2_inv: TensorVariable,
+    sigma: TensorVariable,
+    lmbda_inv: TensorVariable,
+    tau_inv: TensorVariable,
 ) -> Tuple[TensorVariable, TensorVariable]:
-    r"""Gibbs kernel to sample from the posterior distribution of a horsehoe prior.
+    r"""Gibbs kernel to sample from the posterior distribution of a horseshoe prior.
 
     This kernel generates samples from the posterior distribution of the local
-    and global shrinkage parameters of a horsehoe prior, respectively :math:`\lambda`
+    and global shrinkage parameters of a horseshoe prior, respectively :math:`\lambda`
     and :math:`\tau` in the following model:
 
     .. math::
@@ -58,7 +152,7 @@ def horseshoe_step(
         \end{align*}
 
     We use the following observations [1]_ to sample from the posterior
-    conditional probability of :math:`\tau^{-2}` and :math:`\lambda^{-2}`:
+    conditional probability of :math:`\tau` and :math:`\lambda`:
 
     1. The half-Cauchy distribution can be intepreted as a mixture of inverse-gamma
     distributions;
@@ -72,7 +166,7 @@ def horseshoe_step(
         Regression coefficients.
     sigma2
         Variance of the regression coefficients.
-    lambda2_inv
+    lmbda2_inv
         Square inverse of the local shrinkage parameters.
     tau2_inv
         Square inverse of the global shrinkage parameters.
@@ -81,32 +175,85 @@ def horseshoe_step(
     ----------
     ..[1] Makalic, Enes & Schmidt, Daniel. (2016). High-Dimensional Bayesian
           Regularised Regression with the BayesReg Package.
+
     """
-    upsilon_inv = srng.exponential(1 + lambda2_inv)
-    zeta_inv = srng.exponential(1 + tau2_inv)
+    upsilon_inv = srng.exponential(1 + lmbda_inv)
+    zeta_inv = srng.exponential(1 + tau_inv)
 
     beta2 = beta * beta
-    lambda2_inv_new = srng.exponential(upsilon_inv + 0.5 * beta2 * tau2_inv / sigma2)
-    tau2_inv_new = srng.gamma(
+    lmbda_inv_new = srng.exponential(upsilon_inv + 0.5 * beta2 * tau_inv / sigma)
+    tau_inv_new = srng.gamma(
         0.5 * (beta.shape[0] + 1),
-        zeta_inv + 0.5 * (beta2 * lambda2_inv_new).sum() / sigma2,
+        zeta_inv + 0.5 * (beta2 * lmbda_inv_new).sum() / sigma,
     )
-    return lambda2_inv_new, tau2_inv_new
+    return lmbda_inv_new, tau_inv_new
 
 
-def horseshoe_nbinom(
-    srng: at.random.RandomStream,
-    X: TensorVariable,
-    y: TensorVariable,
-    h: TensorVariable,
-    beta_init: TensorVariable,
-    local_shrinkage_init: TensorVariable,
-    global_shrinkage_init: TensorVariable,
-    n_samples: TensorVariable,
-) -> Tuple[List[Variable], Dict]:
-    r"""
-    Build a symbolic graph that describes the gibbs sampler of the negative
-    binomial regression with a HorseShoe prior on the regression coefficients.
+X_lv = var()
+beta_lv = var()
+neg_one_lv = var()
+
+sigmoid_dot_pattern = etuple(
+    etuplize(at.sigmoid),
+    etuple(etuplize(at.mul), neg_one_lv, etuple(etuple(Dot), X_lv, beta_lv)),
+)
+
+
+h_lv = var()
+nbinom_sigmoid_dot_pattern = etuple(
+    etuplize(at.random.nbinom), var(), var(), var(), h_lv, sigmoid_dot_pattern
+)
+
+
+def nbinom_sigmoid_dot_match(
+    graph: TensorVariable,
+) -> Tuple[TensorVariable, TensorVariable, TensorVariable]:
+    graph_opt = optimize_graph(graph)
+    graph_et = etuplize(graph_opt)
+    s = unify(graph_et, nbinom_sigmoid_dot_pattern)
+    if s is False:
+        raise ValueError("Not a negative binomial regression.")
+
+    if all(s[neg_one_lv].data != -1):
+        raise ValueError(
+            "Not a negative binomial regression. The argument to "
+            + "the sigmoid must be minus the dot product."
+        )
+
+    h = eval_if_etuple(s[h_lv])
+    beta_rv = eval_if_etuple(s[beta_lv])
+    X = eval_if_etuple(s[X_lv])
+
+    return X, h, beta_rv
+
+
+def nbinom_horseshoe_model(srng: RandomStream) -> TensorVariable:
+    """Negative binomial regression model with a horseshoe shrinkage prior."""
+    X = at.matrix("X")
+    h = at.scalar("h")
+
+    beta_rv = horseshoe_model(srng)
+    eta = X @ beta_rv
+    p = at.sigmoid(-eta)
+    Y_rv = srng.nbinom(h, p)
+
+    return Y_rv
+
+
+def nbinom_horseshoe_match(
+    Y_rv: TensorVariable,
+) -> Tuple[
+    TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable
+]:
+    X, h, beta_rv = nbinom_sigmoid_dot_match(Y_rv)
+    lmbda_rv, tau_rv = horseshoe_match(beta_rv)
+    return h, X, beta_rv, lmbda_rv, tau_rv
+
+
+def nbinom_horseshoe_gibbs(
+    srng: RandomStream, Y_rv: TensorVariable, y: TensorVariable, num_samples: int
+) -> Tuple[Union[TensorVariable, List[TensorVariable]], Dict]:
+    r"""Build a Gibbs sampler for the negative binomial regression with a horseshoe prior.
 
     The implementation follows the sampler described in [1]. It is designed to
     sample efficiently from the following negative binomial regression model:
@@ -118,8 +265,7 @@ def horseshoe_nbinom(
             h &\sim \pi_h(h) \mathrm{d}h\\
             \pi_i &= \frac{\exp(\psi_i)}{1 + \exp(\psi_i)}\\
             \psi_i &= x^T \beta\\
-            \beta_j &\sim \operatorname{Normal}(0, \lambda_j^2\;\tau^2\;\sigma^2)\\
-            \sigma^2 &\sim \sigma^{-2} \mathrm{d} \sigma\\
+            \beta_j &\sim \operatorname{Normal}(0, \lambda_j^2\;\tau^2)\\
             \lambda_j &\sim \operatorname{HalfCauchy}(0, 1)\\
             \tau &\sim \operatorname{HalfCauchy}(0, 1)
         \end{align*}
@@ -129,30 +275,18 @@ def horseshoe_nbinom(
     ----------
     srng: symbolic random number generator
         The random number generating object to be used during sampling.
-    X: TensorVariable
-        The covariate matrix.
+    Y_rv
+        Model graph.
     y: TensorVariable
         The observed count data.
-    h: TensorVariable
-        The "number of successes" parameter of the negative binomial disribution
-        used to model the data.
-    beta_init: TensorVariable
-        A tensor describing the starting values of the posterior samples of
-        the regression coefficients.
-    local_shrinkage_init: TensorVariable
-        A tensor describing the starting values of the local shrinkage
-        parameter of the horseshoe prior.
-    global_shrinkage_init: TensorVariable
-        A tensor describing the starting values of the global shrinkage
-        parameter of the horseshoe prior.
     n_samples: TensorVariable
         A tensor describing the number of posterior samples to generate.
 
     Returns
     -------
     (outputs, updates): tuple
-        A symbolic sescription of the sampling result to be used when
-        compiling a function to be used for sampling.
+        A symbolic description of the sampling result to be used to
+        compile a sampling function.
 
     Notes
     -----
@@ -172,104 +306,115 @@ def horseshoe_nbinom(
     ..[3] Neelon, Brian. (2019). Bayesian Zero-Inflated Negative Binomial
           Regression Based on Pólya-Gamma Mixtures. Bayesian Anal.
           2019 September ; 14(3): 829–855. doi:10.1214/18-ba1132.
-    """
-    beta = at.as_tensor_variable(beta_init, "beta", dtype=aesara.config.floatX)
-    if beta.ndim != 1:
-        raise ValueError(
-            "`beta_init` must be a vector whose length is equal to the "
-            "number of columns in the covariate matrix"
-        )
-
-    lambda2 = at.as_tensor_variable(
-        local_shrinkage_init, name="lambda2", dtype=aesara.config.floatX
-    )
-    if lambda2.ndim != 1:
-        raise ValueError(
-            "The local shrinkage initial value must be a vector whose size "
-            "is equal to the number of columns in the covariate matrix"
-        )
-
-    tau2 = at.as_tensor_variable(
-        global_shrinkage_init, name="tau2", dtype=aesara.config.floatX
-    )
-    if tau2.ndim != 0:
-        raise ValueError("The global shrinkage initial value must be a scalar")
 
-    n_samples = at.as_tensor_variable(n_samples, name="n_samples", dtype=int)
-    if n_samples.ndim != 0:
-        raise ValueError("The number of samples should be an integer scalar")
-
-    # sigma2 must be set to 1 during sampling for this sampler and the
-    # logistic regression version. It is only required if the data is
-    # modelled using linear regression.
-    sigma2 = at.constant(1, "sigma2")
+    """
 
-    def step_fn(
+    def nbinom_horseshoe_step(
         beta: TensorVariable,
-        lambda2_inv: TensorVariable,
-        tau2_inv: TensorVariable,
-        sigma2: TensorVariable,
-        X: TensorVariable,
+        lmbda: TensorVariable,
+        tau: TensorVariable,
         y: TensorVariable,
+        X: TensorVariable,
         h: TensorVariable,
     ) -> Tuple[TensorVariable, TensorVariable, TensorVariable]:
-        """
-        Complete one full update of the gibbs sampler and return the new state
+        """Complete one full update of the gibbs sampler and return the new state
         of the posterior conditional parameters.
 
         Parameters
         ----------
         beta: Tensorvariable
-            Coefficients of the regression model.
-        lambda2_inv
-            Square inverse of the local shrinkage parameter of the horseshoe prior.
-        tau2_inv
-            Square inverse of the global shrinkage parameters of the horseshoe prior.
-        sigma2
-            Variance of the regression coefficients.
-        X: TensorVariable
-            The covariate matrix.
+            Coefficients (other than intercept) of the regression model.
+        lmbda
+            Inverse of the local shrinkage parameter of the horseshoe prior.
+        tau
+            Inverse of the global shrinkage parameters of the horseshoe prior.
         y: TensorVariable
             The observed count data.
+        X: TensorVariable
+            The covariate matrix.
         h: TensorVariable
             The "number of successes" parameter of the negative binomial disribution
             used to model the data.
+
         """
         xb = X @ beta
         w = srng.gen(polyagamma, y + h, xb)
-
         z = 0.5 * (y - h) / w
-        beta_new = update_beta(srng, w, lambda2_inv * tau2_inv, X, z)
 
-        lambda2_inv_new, tau2_inv_new = horseshoe_step(
-            srng, beta_new, sigma2, lambda2_inv, tau2_inv
+        lmbda_inv = 1.0 / lmbda
+        tau_inv = 1.0 / tau
+        beta_new = update_beta(srng, w, lmbda_inv * tau_inv, X, z)
+
+        lmbda_inv_new, tau_inv_new = horseshoe_step(
+            srng, beta_new, 1.0, lmbda_inv, tau_inv
         )
 
-        return beta_new, lambda2_inv_new, tau2_inv_new
+        return beta_new, 1.0 / lmbda_inv_new, 1.0 / tau_inv_new
+
+    h, X, beta_rv, lmbda_rv, tau_rv = nbinom_horseshoe_match(Y_rv)
 
     outputs, updates = aesara.scan(
-        step_fn,
-        outputs_info=[beta, 1 / lambda2, 1 / tau2],
-        non_sequences=[sigma2, X, y, h],
-        n_steps=n_samples,
+        nbinom_horseshoe_step,
+        outputs_info=[beta_rv, lmbda_rv, tau_rv],
+        non_sequences=[y, X, h],
+        n_steps=num_samples,
         strict=True,
     )
+
     return outputs, updates
 
 
-def horseshoe_logistic(
-    srng: RandomStream,
-    X: TensorVariable,
-    y: TensorVariable,
-    beta_init: TensorVariable,
-    local_shrinkage_init: TensorVariable,
-    global_shrinkage_init: TensorVariable,
-    n_samples: TensorVariable,
-) -> Tuple[List[Variable], Dict]:
-    r"""
+bernoulli_sigmoid_dot_pattern = etuple(
+    etuplize(at.random.bernoulli), var(), var(), var(), sigmoid_dot_pattern
+)
+
+
+def bernoulli_sigmoid_dot_match(
+    graph: TensorVariable,
+) -> Tuple[TensorVariable, TensorVariable]:
+    graph_opt = optimize_graph(graph)
+    graph_et = etuplize(graph_opt)
+    s = unify(graph_et, bernoulli_sigmoid_dot_pattern)
+    if s is False:
+        raise ValueError("Not a Bernoulli regression.")
+
+    if all(s[neg_one_lv].data != -1):
+        raise ValueError(
+            "Not a Bernoulli regression. The argument to the sigmoid "
+            + "must be minus the dot product."
+        )
+
+    beta_rv = eval_if_etuple(s[beta_lv])
+    X = eval_if_etuple(s[X_lv])
+
+    return X, beta_rv
+
+
+def bernoulli_horseshoe_model(srng: RandomStream) -> TensorVariable:
+    """Bernoulli regression model with a horseshoe shrinkage prior."""
+    X = at.matrix("X")
+
+    beta_rv = horseshoe_model(srng)
+    eta = X @ beta_rv
+    p = at.sigmoid(-eta)
+    Y_rv = srng.bernoulli(p)
+
+    return Y_rv
+
+
+def bernoulli_horseshoe_match(
+    Y_rv: TensorVariable,
+) -> Tuple[TensorVariable, TensorVariable, TensorVariable, TensorVariable]:
+    X, beta_rv = bernoulli_sigmoid_dot_match(Y_rv)
+    lmbda_rv, tau_rv = horseshoe_match(beta_rv)
+
+    return X, beta_rv, lmbda_rv, tau_rv
 
-    Build a symbolic graph that describes the gibbs sampler of the binary
-    logistic regression with a Horseshoe prior on the regression coefficients.
+
+def bernoulli_horseshoe_gibbs(
+    srng: RandomStream, Y_rv: TensorVariable, y: TensorVariable, num_samples: int
+) -> Tuple[Union[TensorVariable, List[TensorVariable]], Dict]:
+    r"""Build a Gibbs sampler for bernoulli (logistic) regression with a horseshoe prior.
 
     The implementation follows the sampler described in [1]. It is designed to
     sample efficiently from the following binary logistic regression model:
@@ -277,11 +422,10 @@ def horseshoe_logistic(
     .. math::
 
         \begin{align*}
-            y_i &\sim \operatorname{Benoulli}\left(\pi_i\right)\\
+            y_i &\sim \operatorname{Bernoulli}\left(\pi_i\right)\\
             \pi &= \frac{1}{1 + \exp\left(-(\beta_0 + x_i^T\,\beta)\right)}\\
-            \beta_j &\sim \operatorname{Normal}(0, \lambda_j^2\;\tau^2\;\sigma^2)\\
-            \sigma^2 &\sim \sigma^{-2} \mathrm{d} \sigma\\
-            \lambda_j^2 &\sim \operatorname{HalfCauchy}(0, 1)\\
+            \beta_j &\sim \operatorname{Normal}(0, \lambda_j^2\;\tau^2)\\
+            \lambda_j &\sim \operatorname{HalfCauchy}(0, 1)\\
             \tau &\sim \operatorname{HalfCauchy}(0, 1)
         \end{align*}
 
@@ -289,27 +433,21 @@ def horseshoe_logistic(
     ----------
     srng: symbolic random number generator
         The random number generating object to be used during sampling.
+    Y_rv
+        Model graph.
+    y: TensorVariable
+        The observed binary data.
     X: TensorVariable
         The covariate matrix.
-    y: TensorVariable
-        The observed count data.
-    beta_init: TensorVariable
-        A tensor describing the starting values of the posterior samples of
-        the regression coefficients.
-    local_shrinkage_init: TensorVariable
-        A tensor describing the starting values of the local shrinkage
-        parameter of the horseshoe prior.
-    global_shrinkage_init: TensorVariable
-        A tensor describing the starting values of the global shrinkage
-        parameter of the horseshoe prior.
     n_samples: TensorVariable
         A tensor describing the number of posterior samples to generate.
 
     Returns
     -------
     (outputs, updates): tuple
-        A symbolic sescription of the sampling result to be used when
-        compiling a function to be used for sampling.
+        A symbolic description of the sampling result to be used to
+        compile a sampling function.
+
 
     References
     ----------
@@ -320,80 +458,52 @@ def horseshoe_logistic(
 
     """
 
-    beta = at.as_tensor_variable(beta_init, "beta", dtype=aesara.config.floatX)
-    if beta.ndim != 1:
-        raise ValueError(
-            "`beta_init` must be a vector whose length is equal to the "
-            "number of columns in the covariate matrix"
-        )
-
-    lambda2 = at.as_tensor_variable(
-        local_shrinkage_init, name="lambda2", dtype=aesara.config.floatX
-    )
-    if lambda2.ndim != 1:
-        raise ValueError(
-            "The local shrinkage initial value must be a vector whose size "
-            "is equal to the number of columns in the covariate matrix"
-        )
-
-    tau2 = at.as_tensor_variable(
-        global_shrinkage_init, name="tau2", dtype=aesara.config.floatX
-    )
-    if tau2.ndim != 0:
-        raise ValueError("The global shrinkage initial value must be a scalar")
-
-    n_samples = at.as_tensor_variable(n_samples, name="n_samples", dtype=int)
-    if n_samples.ndim != 0:
-        raise ValueError("The number of samples should be an integer scalar")
-
-    # sigma2 must be set to 1 during sampling for this sampler and the negative
-    # binomial regression version. It is only required if the data is modelled
-    # using linear regression.
-    sigma2 = at.constant(1, "sigma2")
-
-    def step_fn(
+    def bernoulli_horseshoe_step(
         beta: TensorVariable,
-        lambda2_inv: TensorVariable,
-        tau2_inv: TensorVariable,
-        sigma2: TensorVariable,
-        X: TensorVariable,
+        lmbda: TensorVariable,
+        tau: TensorVariable,
         y: TensorVariable,
+        X: TensorVariable,
     ) -> Tuple[TensorVariable, TensorVariable, TensorVariable]:
-        """
-        Complete one full update of the gibbs sampler and return the new state
-        of the posterior conditional parameters.
+        """Complete one full update of the gibbs sampler and return the new
+        state of the posterior conditional parameters.
 
         Parameters
         ----------
-        beta: Tensorvariable
-            Coefficients of the regression model.
-        lambda2_inv
-            Square inverse of the local shrinkage parameter of the horseshoe prior.
-        tau2_inv
-            Square inverse of the global shrinkage parameters of the horseshoe prior.
-        X: TensorVariable
-            The covariate matrix.
+        beta
+            Coefficients (other than intercept) of the regression model.
+        lmbda
+            Square of the local shrinkage parameter of the horseshoe prior.
+        tau
+            Square of the global shrinkage parameters of the horseshoe prior.
         y: TensorVariable
             The observed binary data
-        """
+        X: TensorVariable
+            The covariate matrix.
 
+        """
         xb = X @ beta
         w = srng.gen(polyagamma, 1, xb)
+        z = 0.5 * y / w
 
-        z = (y - 0.5) / w
-        beta_new = update_beta(srng, w, lambda2_inv * tau2_inv, X, z)
+        lmbda_inv = 1.0 / lmbda
+        tau_inv = 1.0 / tau
+        beta_new = update_beta(srng, w, lmbda_inv * tau_inv, X, z)
 
-        lambda2_inv_new, tau2_inv_new = horseshoe_step(
-            srng, beta_new, sigma2, lambda2_inv, tau2_inv
+        lmbda_inv_new, tau_inv_new = horseshoe_step(
+            srng, beta_new, 1.0, lmbda_inv, tau_inv
         )
 
-        return beta_new, lambda2_inv_new, tau2_inv_new
+        return beta_new, 1 / lmbda_inv_new, 1.0 / tau_inv_new
+
+    X, beta_rv, lmbda_rv, tau_rv = bernoulli_horseshoe_match(Y_rv)
 
     outputs, updates = aesara.scan(
-        step_fn,
-        outputs_info=[beta, 1 / lambda2, 1 / tau2],
-        non_sequences=[sigma2, X, y],
-        n_steps=n_samples,
+        bernoulli_horseshoe_step,
+        outputs_info=[beta_rv, lmbda_rv, tau_rv],
+        non_sequences=[y, X],
+        n_steps=num_samples,
         strict=True,
     )
+
     return outputs, updates
diff --git a/tests/test_gibbs.py b/tests/test_gibbs.py
index 734db5b..023429c 100644
--- a/tests/test_gibbs.py
+++ b/tests/test_gibbs.py
@@ -5,147 +5,188 @@
 from aesara.tensor.random.utils import RandomStream
 from scipy.linalg import toeplitz
 
-from aemcmc.gibbs import horseshoe_logistic, horseshoe_nbinom
+from aemcmc.gibbs import (
+    bernoulli_horseshoe_gibbs,
+    bernoulli_horseshoe_match,
+    bernoulli_horseshoe_model,
+    horseshoe_match,
+    horseshoe_model,
+    nbinom_horseshoe_gibbs,
+    nbinom_horseshoe_match,
+    nbinom_horseshoe_model,
+)
 
 
-def test_horseshoe_nbinom():
+@pytest.fixture
+def srng():
+    return RandomStream(1234)
+
+
+def test_match_horseshoe(srng):
+    horseshoe_match(horseshoe_model(srng))
+
+    # exchanging tau and lmbda in the product
+    size = at.scalar("size", dtype="int32")
+    tau_rv = srng.halfcauchy(0, 1, size=1)
+    lmbda_rv = srng.halfcauchy(0, 1, size=size)
+    beta_rv = srng.normal(0, lmbda_rv * tau_rv, size=size)
+    horseshoe_match(beta_rv)
+
+
+def test_match_horseshoe_wrong_graph(srng):
+    beta_rv = srng.normal(0, 1)
+    with pytest.raises(ValueError):
+        horseshoe_match(beta_rv)
+
+
+def test_match_horseshoe_wrong_local_scale_dist(srng):
+    size = at.scalar("size", dtype="int32")
+    tau_rv = srng.halfcauchy(0, 1, size=1)
+    lmbda_rv = srng.normal(0, 1, size=size)
+    beta_rv = srng.normal(0, tau_rv * lmbda_rv, size=size)
+    with pytest.raises(ValueError):
+        horseshoe_match(beta_rv)
+
+
+def test_match_horseshoe_wrong_global_scale_dist(srng):
+    size = at.scalar("size", dtype="int32")
+    tau_rv = srng.normal(0, 1, size=1)
+    lmbda_rv = srng.halfcauchy(0, 1, size=size)
+    beta_rv = srng.normal(0, tau_rv * lmbda_rv, size=size)
+    with pytest.raises(ValueError):
+        horseshoe_match(beta_rv)
+
+
+def test_match_horseshoe_wrong_dimensions(srng):
+    size = at.scalar("size", dtype="int32")
+    tau_rv = srng.halfcauchy(0, 1, size=size)
+    lmbda_rv = srng.halfcauchy(0, 1, size=size)
+
+    beta_rv = srng.normal(0, tau_rv * lmbda_rv, size=size)
+    with pytest.raises(ValueError):
+        horseshoe_match(beta_rv)
+
+
+def test_match_nbinom_horseshoe(srng):
+    nbinom_horseshoe_match(nbinom_horseshoe_model(srng))
+
+
+def test_match_binom_horseshoe_wrong_graph(srng):
+    beta = at.vector("beta")
+    X = at.matrix("X")
+    Y = X @ beta
+
+    with pytest.raises(ValueError):
+        nbinom_horseshoe_match(Y)
+
+
+def test_match_nbinom_horseshoe_wrong_sign(srng):
+    X = at.matrix("X")
+    h = at.scalar("h")
+
+    beta_rv = horseshoe_model(srng)
+    eta = X @ beta_rv
+    p = at.sigmoid(2 * eta)
+    Y_rv = srng.nbinom(h, p)
+
+    with pytest.raises(ValueError):
+        nbinom_horseshoe_match(Y_rv)
+
+
+def test_horseshoe_nbinom(srng):
     """
     This test example is modified from section 3.2 of Makalic & Schmidt (2016)
     """
     h = 2
     p = 10
     N = 50
-    srng = RandomStream(seed=12345)
+
+    # generate synthetic data
     true_beta = np.array([5, 3, 3, 1, 1] + [0] * (p - 5))
     S = toeplitz(0.5 ** np.arange(p))
     X = srng.multivariate_normal(np.zeros(p), cov=S, size=N)
-    y = srng.nbinom(h, at.sigmoid(-X.dot(true_beta)))
-
-    beta_init, lambda2_init, tau2_init, n_samples = (
-        at.vector("beta_init"),
-        at.vector("lambda2_init"),
-        at.scalar("tau2_init"),
-        at.iscalar("n_samples"),
-    )
-    outputs, updates = horseshoe_nbinom(
-        srng, X, y, h, beta_init, lambda2_init, tau2_init, n_samples
-    )
-    sample_fn = aesara.function(
-        [beta_init, lambda2_init, tau2_init, n_samples], outputs, updates=updates
-    )
-
-    rng = np.random.default_rng(54321)
-    posterior_samples = 2000
-    beta, lambda2_inv, tau2_inv = sample_fn(
-        rng.normal(0, 5, size=p), np.ones(p), 1, posterior_samples
-    )
-
-    assert beta.shape == (posterior_samples, p)
-    assert lambda2_inv.shape == (posterior_samples, p)
-    assert tau2_inv.shape == (posterior_samples,)
+    y = srng.nbinom(h, at.sigmoid(-(X.dot(true_beta))))
+
+    # build the model
+    tau_rv = srng.halfcauchy(0, 1, size=1)
+    lambda_rv = srng.halfcauchy(0, 1, size=p)
+    beta_rv = srng.normal(0, tau_rv * lambda_rv, size=p)
+
+    eta_tt = X @ beta_rv
+    p_tt = at.sigmoid(-eta_tt)
+    Y_rv = srng.nbinom(h, p_tt)
+
+    # sample from the posterior distributions
+    num_samples = at.scalar("num_samples", dtype="int32")
+    outputs, updates = nbinom_horseshoe_gibbs(srng, Y_rv, y, num_samples)
+    sample_fn = aesara.function((num_samples,), outputs, updates=updates)
+
+    beta, lmbda, tau = sample_fn(2000)
+
+    assert beta.shape == (2000, p)
+    assert lmbda.shape == (2000, p)
+    assert tau.shape == (2000, 1)
 
     # test distribution domains
-    assert np.all(tau2_inv > 0)
-    assert np.all(lambda2_inv > 0)
-
-    # test if the sampler fails with wrong input
-    with pytest.raises(ValueError, match="`beta_init` must be a vector "):
-        wrong_beta_init = at.tensor3("wrong_beta_init")
-        horseshoe_nbinom(
-            srng, X, y, h, wrong_beta_init, lambda2_init, tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The local shrinkage initial value must be a vector"
-    ):
-        wrong_lambda2_init = at.scalar("wrong_lambda2_init")
-        horseshoe_nbinom(
-            srng, X, y, h, beta_init, wrong_lambda2_init, tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The global shrinkage initial value must be a scalar"
-    ):
-        wrong_tau2_init = at.matrix("wrong_tau2_init")
-        horseshoe_nbinom(
-            srng, X, y, h, beta_init, lambda2_init, wrong_tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The number of samples should be an integer scalar"
-    ):
-        wrong_n_samples = at.vector("wrong_n_samples")
-        horseshoe_nbinom(
-            srng, X, y, h, beta_init, lambda2_init, tau2_init, wrong_n_samples
-        )
-
-
-def test_horseshoe_logistic():
-    """
-    This test example is copied from section 3.2 of Makalic & Schmidt (2016)
-    """
+    assert np.all(tau > 0)
+    assert np.all(lmbda > 0)
+
+
+def test_match_bernoulli_horseshoe(srng):
+    bernoulli_horseshoe_match(bernoulli_horseshoe_model(srng))
+
+
+def test_match_bernoulli_horseshoe_wrong_graph(srng):
+    beta = at.vector("beta")
+    X = at.matrix("X")
+    Y = X @ beta
+
+    with pytest.raises(ValueError):
+        bernoulli_horseshoe_match(Y)
+
+
+def test_match_bernoulli_horseshoe_wrong_sign(srng):
+    X = at.matrix("X")
+
+    beta_rv = horseshoe_model(srng)
+    eta = X @ beta_rv
+    p = at.sigmoid(2 * eta)
+    Y_rv = srng.bernoulli(p)
+
+    with pytest.raises(ValueError):
+        bernoulli_horseshoe_match(Y_rv)
+
+
+def test_bernoulli_horseshoe(srng):
     p = 10
     N = 50
-    srng = RandomStream(seed=12345)
+
+    # generate synthetic data
     true_beta = np.array([5, 3, 3, 1, 1] + [0] * (p - 5))
     S = toeplitz(0.5 ** np.arange(p))
     X = srng.multivariate_normal(np.zeros(p), cov=S, size=N)
     y = srng.bernoulli(at.sigmoid(-X.dot(true_beta)))
 
-    beta_init, lambda2_init, tau2_init, n_samples = (
-        at.vector("beta_init"),
-        at.vector("lambda2_init"),
-        at.scalar("tau2_init"),
-        at.iscalar("n_samples"),
-    )
-    outputs, updates = horseshoe_logistic(
-        srng, X, y, beta_init, lambda2_init, tau2_init, n_samples
-    )
-    sample_fn = aesara.function(
-        [beta_init, lambda2_init, tau2_init, n_samples], outputs, updates=updates
-    )
-
-    rng = np.random.default_rng(54321)
-    posterior_samples = 2000
-    beta, lambda2_inv, tau2_inv = sample_fn(
-        rng.normal(0, 5, size=p), np.ones(p), 1, posterior_samples
-    )
-
-    assert beta.shape == (posterior_samples, p)
-    assert lambda2_inv.shape == (posterior_samples, p)
-    assert tau2_inv.shape == (posterior_samples,)
+    # build the model
+    tau_rv = srng.halfcauchy(0, 1, size=1)
+    lambda_rv = srng.halfcauchy(0, 1, size=p)
+    beta_rv = srng.normal(0, tau_rv * lambda_rv, size=p)
+
+    eta_tt = X @ beta_rv
+    p_tt = at.sigmoid(-eta_tt)
+    Y_rv = srng.bernoulli(p_tt)
+
+    # sample from the posterior distributions
+    num_samples = at.scalar("num_samples", dtype="int32")
+    outputs, updates = bernoulli_horseshoe_gibbs(srng, Y_rv, y, num_samples)
+    sample_fn = aesara.function((num_samples,), outputs, updates=updates)
+
+    beta, lmbda, tau = sample_fn(2000)
+
+    assert beta.shape == (2000, p)
+    assert lmbda.shape == (2000, p)
+    assert tau.shape == (2000, 1)
 
     # test distribution domains
-    assert np.all(tau2_inv > 0)
-    assert np.all(lambda2_inv > 0)
-
-    # test if the sampler fails with wrong input
-    with pytest.raises(ValueError, match="`beta_init` must be a vector "):
-        wrong_beta_init = at.tensor3("wrong_beta_init")
-        horseshoe_logistic(
-            srng, X, y, wrong_beta_init, lambda2_init, tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The local shrinkage initial value must be a vector"
-    ):
-        wrong_lambda2_init = at.scalar("wrong_lambda2_init")
-        horseshoe_logistic(
-            srng, X, y, beta_init, wrong_lambda2_init, tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The global shrinkage initial value must be a scalar"
-    ):
-        wrong_tau2_init = at.matrix("wrong_tau2_init")
-        horseshoe_logistic(
-            srng, X, y, beta_init, lambda2_init, wrong_tau2_init, n_samples
-        )
-
-    with pytest.raises(
-        ValueError, match="The number of samples should be an integer scalar"
-    ):
-        wrong_n_samples = at.vector("wrong_n_samples")
-        horseshoe_logistic(
-            srng, X, y, beta_init, lambda2_init, tau2_init, wrong_n_samples
-        )
+    assert np.all(tau > 0)
+    assert np.all(lmbda > 0)