<commit_before># Source code modified from scipy.stats._multivariate.py
#
# Copyright (c) 2001, 2002 Enthought, Inc.
# All rights reserved.
#
# Copyright (c) 2003-2019 SciPy Developers.
# All rights reserved.

import jax.numpy as np
from jax.scipy.special import digamma, gammaln

from numpyro.distributions.util import xlogy, standard_gamma
from numpyro.distributions.distribution import jax_mvcontinuous


def _lnB(alpha):
    return np.sum(gammaln(alpha), axis=-1) - gammaln(np.sum(alpha, axis=-1))


class dirichlet_gen(jax_mvcontinuous):
    # TODO: use dirichlet doc instead of the default one of rv_continuous
    # TODO: add _argcheck, _support_mask with simplex
    def _logpdf(self, x, alpha):
        lnB = _lnB(alpha)
        return -lnB + np.sum(xlogy(alpha - 1, x), axis=-1)

    def _pdf(self, x, alpha):
        return np.exp(self._logpdf(x, alpha))

    def mean(self, alpha):
        return alpha / np.sum(alpha, axis=-1, keepdims=True)

    def var(self, alpha):
        alpha0 = np.sum(alpha, axis=-1, keepdims=True)
        return (alpha * (alpha0 - alpha)) / ((alpha0 * alpha0) * (alpha0 + 1))

    def _entropy(self, alpha):
        alpha0 = np.sum(alpha, axis=-1)
        lnB = _lnB(alpha)
        K = alpha.shape[-1]
        return lnB + (alpha0 - K) * digamma(alpha0) - np.inner((alpha - 1) * digamma(alpha))

    def _rvs(self, alpha):
        K = alpha.shape[-1]
        gamma_samples = standard_gamma(self._random_state, alpha, self._size + (K,))
        return gamma_samples / np.sum(gamma_samples, axis=-1, keepdims=True)


dirichlet = dirichlet_gen(a=0.0, b=1.0, name='dirichlet')<commit_msg>add shape to standard_gamma to avoid recompiling<commit_after># Source code modified from scipy.stats._multivariate.py
#
# Copyright (c) 2001, 2002 Enthought, Inc.
# All rights reserved.
#
# Copyright (c) 2003-2019 SciPy Developers.
# All rights reserved.

import jax.numpy as np
from jax.scipy.special import digamma, gammaln

from numpyro.distributions.util import xlogy, standard_gamma
from numpyro.distributions.distribution import jax_mvcontinuous


def _lnB(alpha):
    return np.sum(gammaln(alpha), axis=-1) - gammaln(np.sum(alpha, axis=-1))


class dirichlet_gen(jax_mvcontinuous):
    # TODO: use dirichlet doc instead of the default one of rv_continuous
    # TODO: add _argcheck, _support_mask with simplex
    def _logpdf(self, x, alpha):
        lnB = _lnB(alpha)
        return -lnB + np.sum(xlogy(alpha - 1, x), axis=-1)

    def _pdf(self, x, alpha):
        return np.exp(self._logpdf(x, alpha))

    def mean(self, alpha):
        return alpha / np.sum(alpha, axis=-1, keepdims=True)

    def var(self, alpha):
        alpha0 = np.sum(alpha, axis=-1, keepdims=True)
        return (alpha * (alpha0 - alpha)) / ((alpha0 * alpha0) * (alpha0 + 1))

    def _entropy(self, alpha):
        alpha0 = np.sum(alpha, axis=-1)
        lnB = _lnB(alpha)
        K = alpha.shape[-1]
        return lnB + (alpha0 - K) * digamma(alpha0) - np.inner((alpha - 1) * digamma(alpha))

    def _rvs(self, alpha):
        K = alpha.shape[-1]
        gamma_samples = standard_gamma(self._random_state, alpha, shape=self._size + (K,))
        return gamma_samples / np.sum(gamma_samples, axis=-1, keepdims=True)


dirichlet = dirichlet_gen(a=0.0, b=1.0, name='dirichlet')
{'fraction_non_alphanumeric': 0.12921860747947705, 'fraction_numerical': 0.024323502584372148, 'mean_word_length': 3.2671854734111543, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3815789', 'n_tokens_mistral': 1219, 'n_tokens_neox': 1136, 'n_words': 297}