Stochastic matrix in pymc3

stochastic_matrix.py

import numpy as np
import pymc3 as pm
import theano

class StochasticMatrix(pm.Continuous):
    """A stochastic matrix has rows that sum to 1."""
    def __init__(self, theta, *args, **kwargs):
        shape = (theta.shape[-1], theta.shape[-1])
        kwargs.setdefault('shape', shape)
        super(StochasticMatrix, self).__init__(*args, **kwargs)
        self.theta = theta
        self.row_distribution = pm.Dirichlet.dist(a=self.theta, transform=None)

    def logp(self, value):
        results, updates = theano.scan(self.row_distribution.logp, value)
        return results.sum()

    def random(self):
        return pm.Dirichlet('rand', self.theta).random(size=len(self.theta))


with Model() as model:
    transition = StochasticMatrix('transition', theta=np.array([1, 1]), testval=np.array([[1, 0], [1, 0]]))
    trace = pm.sample(5)