def __init__(self, K, n, C=None, sigma_C=1, mu=None, mu_pi=None):
        """
        Create a PGMultinomial distribution with mean and covariance for psi.

        :param K:       Dimensionality of the multinomial distribution
        :param mu_C:    Mean of the matrix normal distribution over C
        """
        assert isinstance(K, int) and K >= 2, "K must be an integer >= 2"
        self.K = K

        assert isinstance(n, int) and n >= 1, "n must be an integer >= 1"
        self.n = n

        # Initialize emission matrix C
        self.sigma_C = sigma_C
        if C is None:
            self.C = self.sigma_C * np.random.randn(self.K-1, self.n)
            # mu, sigma = compute_psi_cmoments(np.ones(K))
            # self.C = compute_psi_cmoments(np.ones(K))[0][:,None] * np.ones((self.K-1, self.n))
        else:
            assert C.shape == (self.K-1, self.n)
            self.C = C

        # Initialize the observation mean (mu)
        if mu == None and mu_pi is None:
            self.mu = np.zeros(self.K-1)
        elif mu is not None:
            assert mu.shape == (self.K-1,)
            self.mu = mu
        else:
            assert mu_pi.shape == (self.K,)
            self.mu = pi_to_psi(mu_pi)

        # Initialize Polya-gamma augmentation variables
        self.ppgs = initialize_polya_gamma_samplers()
dataset/ETHPy150Open HIPS/pgmult/pgmult/distributions.py/PGMultinomialRegression.__init__/operator `is` -> `==`