<commit_before>import tensorflow as tf
import tensorflow.contrib.rnn as rnn
import tensorflow.contrib.layers as layers

from vae import vae
from transformer import transformer
from gumbel import gumbel_softmax_binary


class AIRModel:

    def __init__(self, input_images, target_num_digits,
                 max_steps=3, max_digits=2, lstm_units=256, canvas_size=50, windows_size=28,
                 vae_latent_dimensions=50, vae_recognition_units=(512, 256), vae_generative_units=(256, 512),
                 scale_prior_mean=-1.0, scale_prior_variance=0.1, shift_prior_mean=0.0, shift_prior_variance=1.0,
                 vae_prior_mean=0.0, vae_prior_variance=1.0, vae_likelihood_std=0.3,
                 z_pres_prior=0.01, gumbel_temperature=1.0, learning_rate=1e-4, gradient_clipping_norm=10.0,
                 num_summary_images=12, train=False, reuse=False, scope="air",
                 annealing_schedules=None):

        self.input_images = input_images
        self.target_num_digits = target_num_digits
        self.batch_size = tf.shape(input_images)[0]

        self.max_steps = max_steps
        self.max_digits = max_digits
        self.lstm_units = lstm_units
        self.canvas_size = canvas_size
        self.windows_size = windows_size

        self.vae_latent_dimensions = vae_latent_dimensions
        self.vae_recognition_units = vae_recognition_units
        self.vae_generative_units = vae_generative_units

        self.scale_prior_mean = scale_prior_mean
        self.scale_prior_variance = scale_prior_variance
        self.shift_prior_mean = shift_prior_mean
        self.shift_prior_variance = shift_prior_variance
        self.vae_prior_mean = vae_prior_mean
        self.vae_prior_variance = vae_prior_variance
        self.vae_likelihood_std = vae_likelihood_std

        self.z_pres_prior = z_pres_prior
        self.gumbel_temperature = gumbel_temperature
        self.learning_rate = learning_rate
        self.gradient_clipping_norm = gradient_clipping_norm
        self.num_summary_images = num_summary_images

        self.train = train
        self.num_summaries = []
        self.img_summaries = []
        self.var_summaries = []
        self.grad_summaries = []

        with tf.variable_scope(scope, reuse=reuse):
            self.global_step = tf.get_variable(name="global_step", shape=[], dtype=tf.int32,
                                               initializer=tf.constant_initializer(0), trainable=False)

            self.scale_prior_log_variance = tf.log(scale_prior_variance, name="scale_prior_log_variance")
            self.shift_prior_log_variance = tf.log(shift_prior_variance, name="shift_prior_log_variance")
            self.vae_prior_log_variance = tf.log(vae_prior_variance, name="vae_prior_log_variance")

            if annealing_schedules is not None:
                for param, schedule in annealing_schedules.items():
                    setattr(self, param, self._create_annealed_tensor(
                        param, schedule, self.global_step
                    ))

            self.rec_num_digits = None
            self.rec_scales = None
            self.rec_shifts = None
            self.reconstruction = None
            self.loss = None
            self.accuracy = None
            self.training = None

            self._create_model()

    @staticmethod
    def _create_annealed_tensor(param, schedule, global_step):
        return tf.maximum(
            tf.train.exponential_decay(
                schedule["init"], global_step,
                schedule["iters"], schedule["factor"],
                staircase=True, name=param
            ),
            schedule["min"],
            name=param + "_max"
        )

    @staticmethod
    def _sample_from_mvn(mean, diag_variance):
        # sampling from the multivariate normal
        # with given mean and diagonal covaraince
        standard_normal = tf.random_normal(tf.shape(mean))
        return mean + standard_normal * tf.sqrt(diag_variance)

    @staticmethod
    def _draw_colored_bounding_boxes(images, boxes, steps):
        channels = [images, images, images]

        for s in range(3):
            # empty canvas with s-th bounding box
            step_box = tf.expand_dims(boxes[:, s, :, :], 3)

            for c in range(3):
                if s == c:
                    # adding the box to c-th channel
                    # if the number of attention steps is greater than s
                    channels[c] = tf.where(
                        tf.greater(steps, s),
                        tf.minimum(channels[c] + step_box, tf.ones_like(images)),
                        channels[c]
                    )
                else:
                    # subtracting the box from channels other than c-th
                    # if the number of attention steps is greater than s
                    channels[c] = tf.where(
                        tf.greater(steps, s),
                        tf.maximum(channels[c] - step_box, tf.zeros_like(images)),
                        channels[c]
                    )

        # concatenating all three channels to obtain
        # potentially three R, G, and B bounding boxes
        return tf.concat(channels, axis=3)

    def _summarize_by_digit_count(self, tensor, digits, name):
        # converting to float in case of int tensors
        float_tensor = tf.cast(tensor, tf.float32)

        for i in range(self.max_digits+1):
            # summarizing the scalar for only those
            # images that have exactly i digits
            self.num_summaries.append(
                tf.summary.scalar(
                    name + "_" + str(i) + "_dig",
                    tf.reduce_mean(tf.boolean_mask(
                        float_tensor, tf.equal(digits, i)
                    ))
                )
            )

        # summarizing the scalar for all images
        self.num_summaries.append(
            tf.summary.scalar(
                name + "_all_dig",
                tf.reduce_mean(float_tensor)
            )
        )

    def _summarize_by_step(self, tensor, steps, name, one_more_step=False, all_steps=False):
        try:
            for i in range(self.max_steps):
                if all_steps:
                    # summarizing the entire (i+1)-st step without
                    # differentiating between actual step counts
                    self._summarize_by_digit_count(
                        tensor[:, i], self.target_num_digits,
                        name + "_" + str(i+1) + "_step"
                    )
                else:
                    # considering one more step if required by one_more_step=True
                    # by subtracting 1 from loop variable i (e.g. 0 steps > -1)
                    mask = tf.greater(steps, i - (1 if one_more_step else 0))

                    # summarizing (i+1)-st step only for those
                    # batch items that actually had (i+1)-st step
                    self._summarize_by_digit_count(
                        tf.boolean_mask(tensor[:, i], mask),
                        tf.boolean_mask(self.target_num_digits, mask),
                        name + "_" + str(i+1) + "_step"
                    )
        except Exception as ex:
            print("Warning:", ex)

    def _visualize_reconstructions(self, original, reconstruction, st_back, steps, zoom):
        # enlarging the original images
        large_original = tf.image.resize_images(
            tf.reshape(original, [-1, self.canvas_size, self.canvas_size, 1]),
            [zoom * self.canvas_size, zoom * self.canvas_size]
        )

        # enlarging the reconstructions
        large_reconstruction = tf.image.resize_images(
            tf.reshape(reconstruction, [-1, self.canvas_size, self.canvas_size, 1]),
            [zoom * self.canvas_size, zoom * self.canvas_size]
        )

        # drawing the attention windows
        # using backward ST matrices
        num_images = tf.shape(original)[0]
        boxes = tf.reshape(
            tf.clip_by_value(
                transformer(
                    tf.expand_dims(
                        tf.image.draw_bounding_boxes(
                            tf.zeros(
                                [num_images * self.max_steps, self.windows_size, self.windows_size, 1],
                                dtype=reconstruction.dtype
                            ),<commit_msg>The bug causing "out of range" errors while generating by-step summaries has been fixed.<commit_after>import tensorflow as tf
import tensorflow.contrib.rnn as rnn
import tensorflow.contrib.layers as layers

from vae import vae
from transformer import transformer
from gumbel import gumbel_softmax_binary


class AIRModel:

    def __init__(self, input_images, target_num_digits,
                 max_steps=3, max_digits=2, lstm_units=256, canvas_size=50, windows_size=28,
                 vae_latent_dimensions=50, vae_recognition_units=(512, 256), vae_generative_units=(256, 512),
                 scale_prior_mean=-1.0, scale_prior_variance=0.1, shift_prior_mean=0.0, shift_prior_variance=1.0,
                 vae_prior_mean=0.0, vae_prior_variance=1.0, vae_likelihood_std=0.3,
                 z_pres_prior=0.01, gumbel_temperature=1.0, learning_rate=1e-4, gradient_clipping_norm=10.0,
                 num_summary_images=12, train=False, reuse=False, scope="air",
                 annealing_schedules=None):

        self.input_images = input_images
        self.target_num_digits = target_num_digits
        self.batch_size = tf.shape(input_images)[0]

        self.max_steps = max_steps
        self.max_digits = max_digits
        self.lstm_units = lstm_units
        self.canvas_size = canvas_size
        self.windows_size = windows_size

        self.vae_latent_dimensions = vae_latent_dimensions
        self.vae_recognition_units = vae_recognition_units
        self.vae_generative_units = vae_generative_units

        self.scale_prior_mean = scale_prior_mean
        self.scale_prior_variance = scale_prior_variance
        self.shift_prior_mean = shift_prior_mean
        self.shift_prior_variance = shift_prior_variance
        self.vae_prior_mean = vae_prior_mean
        self.vae_prior_variance = vae_prior_variance
        self.vae_likelihood_std = vae_likelihood_std

        self.z_pres_prior = z_pres_prior
        self.gumbel_temperature = gumbel_temperature
        self.learning_rate = learning_rate
        self.gradient_clipping_norm = gradient_clipping_norm
        self.num_summary_images = num_summary_images

        self.train = train
        self.num_summaries = []
        self.img_summaries = []
        self.var_summaries = []
        self.grad_summaries = []

        with tf.variable_scope(scope, reuse=reuse):
            self.global_step = tf.get_variable(name="global_step", shape=[], dtype=tf.int32,
                                               initializer=tf.constant_initializer(0), trainable=False)

            self.scale_prior_log_variance = tf.log(scale_prior_variance, name="scale_prior_log_variance")
            self.shift_prior_log_variance = tf.log(shift_prior_variance, name="shift_prior_log_variance")
            self.vae_prior_log_variance = tf.log(vae_prior_variance, name="vae_prior_log_variance")

            if annealing_schedules is not None:
                for param, schedule in annealing_schedules.items():
                    setattr(self, param, self._create_annealed_tensor(
                        param, schedule, self.global_step
                    ))

            self.rec_num_digits = None
            self.rec_scales = None
            self.rec_shifts = None
            self.reconstruction = None
            self.loss = None
            self.accuracy = None
            self.training = None

            self._create_model()

    @staticmethod
    def _create_annealed_tensor(param, schedule, global_step):
        return tf.maximum(
            tf.train.exponential_decay(
                schedule["init"], global_step,
                schedule["iters"], schedule["factor"],
                staircase=True, name=param
            ),
            schedule["min"],
            name=param + "_max"
        )

    @staticmethod
    def _sample_from_mvn(mean, diag_variance):
        # sampling from the multivariate normal
        # with given mean and diagonal covaraince
        standard_normal = tf.random_normal(tf.shape(mean))
        return mean + standard_normal * tf.sqrt(diag_variance)

    @staticmethod
    def _draw_colored_bounding_boxes(images, boxes, steps):
        channels = [images, images, images]

        for s in range(3):
            # empty canvas with s-th bounding box
            step_box = tf.expand_dims(boxes[:, s, :, :], 3)

            for c in range(3):
                if s == c:
                    # adding the box to c-th channel
                    # if the number of attention steps is greater than s
                    channels[c] = tf.where(
                        tf.greater(steps, s),
                        tf.minimum(channels[c] + step_box, tf.ones_like(images)),
                        channels[c]
                    )
                else:
                    # subtracting the box from channels other than c-th
                    # if the number of attention steps is greater than s
                    channels[c] = tf.where(
                        tf.greater(steps, s),
                        tf.maximum(channels[c] - step_box, tf.zeros_like(images)),
                        channels[c]
                    )

        # concatenating all three channels to obtain
        # potentially three R, G, and B bounding boxes
        return tf.concat(channels, axis=3)

    def _summarize_by_digit_count(self, tensor, digits, name):
        # converting to float in case of int tensors
        float_tensor = tf.cast(tensor, tf.float32)

        for i in range(self.max_digits+1):
            # summarizing the scalar for only those
            # images that have exactly i digits
            self.num_summaries.append(
                tf.summary.scalar(
                    name + "_" + str(i) + "_dig",
                    tf.reduce_mean(tf.boolean_mask(
                        float_tensor, tf.equal(digits, i)
                    ))
                )
            )

        # summarizing the scalar for all images
        self.num_summaries.append(
            tf.summary.scalar(
                name + "_all_dig",
                tf.reduce_mean(float_tensor)
            )
        )

    def _summarize_by_step(self, tensor, steps, name, one_more_step=False, all_steps=False):
        # padding (if required) the number of rows in the tensor
        # up to self.max_steps to avoid possible "out of range" errors
        tensor = tf.pad(tensor, [[0, 0], [0, self.max_steps - tf.shape(tensor)[1]]])

        for i in range(self.max_steps):
            if all_steps:
                # summarizing the entire (i+1)-st step without
                # differentiating between actual step counts
                self._summarize_by_digit_count(
                    tensor[:, i], self.target_num_digits,
                    name + "_" + str(i+1) + "_step"
                )
            else:
                # considering one more step if required by one_more_step=True
                # by subtracting 1 from loop variable i (e.g. 0 steps > -1)
                mask = tf.greater(steps, i - (1 if one_more_step else 0))

                # summarizing (i+1)-st step only for those
                # batch items that actually had (i+1)-st step
                self._summarize_by_digit_count(
                    tf.boolean_mask(tensor[:, i], mask),
                    tf.boolean_mask(self.target_num_digits, mask),
                    name + "_" + str(i+1) + "_step"
                )

    def _visualize_reconstructions(self, original, reconstruction, st_back, steps, zoom):
        # enlarging the original images
        large_original = tf.image.resize_images(
            tf.reshape(original, [-1, self.canvas_size, self.canvas_size, 1]),
            [zoom * self.canvas_size, zoom * self.canvas_size]
        )

        # enlarging the reconstructions
        large_reconstruction = tf.image.resize_images(
            tf.reshape(reconstruction, [-1, self.canvas_size, self.canvas_size, 1]),
            [zoom * self.canvas_size, zoom * self.canvas_size]
        )

        # drawing the attention windows
        # using backward ST matrices
        num_images = tf.shape(original)[0]
        boxes = tf.reshape(
            tf.clip_by_value(
                transformer(
                    tf.expand_dims(
                        tf.image.draw_bounding_boxes(
                            tf.zeros(
                                [num_images * self.max_steps, self.windows_size, self.windows_size, 1],
                                dtype=reconstruction.dtype
                            ),
{'fraction_non_alphanumeric': 0.07983441750443525, 'fraction_numerical': 0.009225310467179185, 'mean_word_length': 1.8213213213213213, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '2699064', 'n_tokens_mistral': 5023, 'n_tokens_neox': 4642, 'n_words': 1017}