<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from copy import copy
from collections import namedtuple

import six
from six.moves import zip
import tensorflow as tf

from zhusuan.utils import merge_dicts


__all__ = [
    "SGMCMC",
    "SGLD",
    "PSGLD",
    "SGHMC",
    "SGNHT",
]


class SGMCMC:
    """
    Base class for stochastic gradient MCMC algorithms.
    """
    def __init__(self):
        self.t = tf.Variable(0, name="t", trainable=False, dtype=tf.int32)

    def sample(self, meta_model, observed, latent, record_full=None, split_k=None, batch_size=None):
        if callable(meta_model):
            # TODO: raise warning
            self._meta_model = None
            self._log_joint = meta_model
        else:
            self._meta_model = meta_model
            self._log_joint = lambda obs: meta_model.observe(**obs).log_joint()

        self._observed = observed
        self._latent = latent

        latent_k, latent_v = [list(i) for i in zip(*six.iteritems(latent))]
        for i, v in enumerate(latent_v):
            if not isinstance(v, tf.Variable):
                raise TypeError("latent['{}'] is not a tensorflow Variable."
                                .format(latent_k[i]))
        qs = copy(latent_v)
        self._define_variables(qs)

        def _get_log_posterior(var_list, observed):
            joint_obs = merge_dicts(dict(zip(latent_k, var_list)), observed)
            return self._log_joint(joint_obs)

        def _get_gradient_direct(var_list, observed):
            return tf.gradients(_get_log_posterior(var_list, observed), var_list)

        def _get_gradient(var_list):
            if split_k is None or len(split_k) == 0:
                return _get_gradient_direct(var_list, observed)
            else:
                obs_size = tf.shape(observed[split_k[0]])[0]
                iters = (obs_size-1) // batch_size + 1
                def _get_gradient_sum(var_list):
                    other = copy(observed)
                    for k in split_k:
                        del other[k]
                    slice_v = [observed[k] for k in split_k]
                    obs_v = [tf.Variable(lambda: tf.zeros([batch_size] + list(v.shape)[1:], v.dtype)) for v in slice_v]
                    obs = merge_dicts(dict(zip(split_k, obs_v)), other)
                    gradient = _get_gradient_direct(var_list, obs)

                    def cond(i, total_gradient):
                        return tf.less(i, iters)

                    def body(i, total_gradient):
                        with tf.control_dependencies([ov.assign(v[i * batch_size:(i + 1) * batch_size])
                                for (ov, v) in zip(obs_v, slice_v)]):
                            total_gradient = [old + this / tf.cast(iters, tf.float32)
                                for (old, this) in zip(total_gradient, gradient)]
                        return i+1, total_gradient

                    i = tf.constant(0)
                    gradients = [tf.zeros_like(grad) for grad in _get_gradient_direct(var_list, observed)]

                    return tf.while_loop(cond, body, [i, gradients])[1]

                return tf.cond(tf.greater(obs_size, batch_size), lambda: _get_gradient_sum(var_list),
                               lambda: _get_gradient_direct(var_list, observed))

        if record_full is not None:
            self._old_values = [tf.Variable(tf.zeros_like(q)) for q in qs]
            self._full_gradients = [tf.Variable(tf.zeros_like(grad)) for grad in _get_gradient_direct(qs, observed)]

        def get_gradient(var_list):
            if record_full is None:
                return _get_gradient(var_list)
            else:
                def collect_full_gradients():
                    # Remark: Calculating full gradients need a large memory, so we can calculate batchly and average them.
                    assign_grads = [full_grad.assign(grad) for (full_grad, grad) in
                        zip(self._full_gradients, _get_gradient(var_list))]
                    assign_params = [old.assign(new) for (old, new) in zip(self._old_values, var_list)]
                    with tf.control_dependencies([tf.group(*assign_grads), tf.group(*assign_params)]):
                        return [tf.identity(grad) for grad in self._full_gradients]

                def estimate_gradients():
                    # Remark: Objective function should not change except for random batching
                    now_batch_gradients = _get_gradient(var_list)
                    old_batch_gradients = _get_gradient(self._old_values)<commit_msg>Fix some bugs and problems regarding splitting in sgmcmc.py.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from copy import copy
from collections import namedtuple

import six
from six.moves import zip
import tensorflow as tf

from zhusuan.utils import merge_dicts


__all__ = [
    "SGMCMC",
    "SGLD",
    "PSGLD",
    "SGHMC",
    "SGNHT",
]


class SGMCMC:
    """
    Base class for stochastic gradient MCMC algorithms.
    """
    def __init__(self):
        self.t = tf.Variable(0, name="t", trainable=False, dtype=tf.int32)

    def sample(self, meta_model, observed, latent, record_full=None, batch_size=None, split_k=None):
        # TODO: remove direct; use tf.greater
        if callable(meta_model):
            # TODO: raise warning
            self._meta_model = None
            self._log_joint = meta_model
        else:
            self._meta_model = meta_model
            self._log_joint = lambda obs: meta_model.observe(**obs).log_joint()

        self._observed = observed
        self._latent = latent
        if batch_size is not None and split_k is None:
            split_k = list(observed.keys())

        latent_k, latent_v = [list(i) for i in zip(*six.iteritems(latent))]
        for i, v in enumerate(latent_v):
            if not isinstance(v, tf.Variable):
                raise TypeError("latent['{}'] is not a tensorflow Variable."
                                .format(latent_k[i]))
        qs = copy(latent_v)
        self._define_variables(qs)

        def _get_log_posterior(var_list, observed):
            joint_obs = merge_dicts(dict(zip(latent_k, var_list)), observed)
            return self._log_joint(joint_obs)

        def _get_gradient_direct(var_list, observed):
            return tf.gradients(_get_log_posterior(var_list, observed), var_list)

        def _get_gradient(var_list, direct=False):
            if batch_size is None or len(split_k) == 0 or direct:
                return _get_gradient_direct(var_list, observed)
            else:
                obs_size = tf.shape(observed[split_k[0]])[0]
                iters = (obs_size-1) // batch_size + 1
                def _get_gradient_sum(var_list):
                    other = copy(observed)
                    for k in split_k:
                        del other[k]
                    slice_v = [tf.convert_to_tensor(observed[k]) for k in split_k]

                    def cond(i, total_gradient):
                        return tf.less(i, iters)

                    def body(i, total_gradient):
                        assigned = [v[i * batch_size:(i + 1) * batch_size] for v in slice_v]
                        obs_body = merge_dicts(dict(zip(split_k, assigned)), other)
                        true_batch_size = tf.shape(observed[split_k[0]][i * batch_size:(i + 1) * batch_size])[0]
                        discount = tf.cast(true_batch_size, tf.float32) / tf.cast(obs_size, tf.float32)
                        total_gradient = [old + this * discount
                            for (old, this) in zip(total_gradient, _get_gradient_direct(var_list, obs_body))]
                        return i+1, total_gradient

                    i = tf.constant(0)
                    gradients = [tf.zeros_like(grad) for grad in _get_gradient_direct(var_list, observed)]

                    return tf.while_loop(cond, body, [i, gradients], parallel_iterations=1)[1]

                return _get_gradient_sum(var_list)
                # return tf.cond(tf.greater(obs_size, batch_size), lambda: _get_gradient_sum(var_list),
                               # lambda: _get_gradient_direct(var_list, observed))

        if record_full is not None:
            self._old_values = [tf.Variable(tf.zeros_like(q)) for q in qs]
            self._full_gradients = [tf.Variable(tf.zeros_like(grad)) for grad in _get_gradient_direct(qs, observed)]

        def get_gradient(var_list):
            if record_full is None:
                return _get_gradient(var_list)
            else:
                def collect_full_gradients():
                    # Remark: Calculating full gradients need a large memory, so we can calculate batchly and average them.
                    assign_grads = [full_grad.assign(grad) for (full_grad, grad) in
                        zip(self._full_gradients, _get_gradient(var_list))]
                    assign_params = [old.assign(new) for (old, new) in zip(self._old_values, var_list)]
                    with tf.control_dependencies([tf.group(*assign_grads), tf.group(*assign_params)]):
                        return [tf.identity(grad) for grad in self._full_gradients]

                def estimate_gradients():
                    # Remark: Objective function should not change except for random batching
                    now_batch_gradients = _get_gradient(var_list, direct=True)
                    old_batch_gradients = _get_gradient(self._old_values, direct=True)
{'fraction_non_alphanumeric': 0.09281870860927152, 'fraction_numerical': 0.0038286423841059603, 'mean_word_length': 1.995041834521227, '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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3467395', 'n_tokens_mistral': 2953, 'n_tokens_neox': 2743, 'n_words': 667}