def _per_timestep_call(self, current_outputs, new_inputs, length, timestep, **kwargs): """Returns Tensor of shape [..., timestep+1, vocab_size]. Args: current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far generated sequence Tensor. new_inputs: Tensor of shape [..., vocab_size], the new input to generate its output given current_outputs. length: Length of final desired sequence. timestep: Current timestep. **kwargs: Optional keyword arguments to layer. """ inputs = tf.concat([current_outputs, new_inputs[..., tf.newaxis, :]], axis=-2) batch_ndims = inputs.shape.ndims - 2 padded_inputs = tf.pad(inputs, paddings=[[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]]) net = self.layer(padded_inputs, **kwargs) if net.shape[-1] == self.vocab_size: loc = net loc = loc[..., :(timestep + 1), :] loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) new_outputs = utils.one_hot_minus(inputs, loc) else: raise ValueError( 'Output of layer does not have compatible dimensions.') outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2) if not tf.executing_eagerly(): outputs.set_shape([None] * batch_ndims + [timestep + 1, self.vocab_size]) return outputs
def reverse(self, inputs, **kwargs): """Reverse pass returning the inverse autoregressive transformation.""" if not self.built: self._maybe_build(inputs) net = self.layer(inputs, **kwargs) if net.shape[-1] == self.subset_K: #! loc = net else: raise ValueError( 'Output of layer does not have compatible dimensions.') loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) # operate on subset: x = inputs x = tf.gather(x, self.shuffling, axis=-1) x1 = x[..., :self.subset_K] #x1 = super().call(x1) x1 = utils.one_hot_add(loc, x1) x2 = x[..., self.subset_K:] x = tf.concat([x1, x2], axis=-1) x = tf.gather(x, self.inverted_shuffling, axis=-1) outputs = x return outputs
def _initial_call(self, new_inputs, length, **kwargs): inputs = new_inputs[..., tf.newaxis, :] batch_ndims = inputs.shape.ndims - 2 padded_inputs = tf.pad(inputs, paddings=[[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]]) net = self.layer(padded_inputs, **kwargs) if net.shape[-1] == self.subset_K: #! loc = net loc = loc[..., 0:1, :] loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) # operate on subset: x = inputs x = tf.gather(x, self.shuffling, axis=-1) x1 = x[..., :self.subset_K] #x1 = super().call(x1) x1 = utils.one_hot_minus( x1, loc) #outputs = utils.one_hot_minus(inputs, loc) x2 = x[..., self.subset_K:] x = tf.concat([x1, x2], axis=-1) x = tf.gather(x, self.inverted_shuffling, axis=-1) outputs = x else: raise ValueError( 'Output of layer does not have compatible dimensions.') return outputs
def log_prob(self, sample, eps_prob=1e-31): B = self.components.probs.shape[-2] component_probs = self.components.probs shift = utils.one_hot_argmax(self.logits, self.temperature) sample = utils.one_hot_add(sample[:, :, None, :], shift[None, :, :, :]) prob = tf.reduce_sum(component_probs * sample + eps_prob, -1) # sum over categories => n x N x B log_prob = tf.math.log(prob) + np.log(1. / B) # n x N x B log_prob = tf.math.reduce_logsumexp( log_prob, -1) # sum over B mixture components => n x N return tf.reduce_sum(log_prob, -1) # sum over N
def sample_extm(self, n=1, *args, **kwargs): B = self.components.probs.shape[-2] dtype = self.dtype shift = utils.one_hot_argmax(self.logits, self.temperature) # N x B x K sample = self.components.sample(n) # n x N x B x K sample = utils.one_hot_minus(sample, shift) selected_flows = (np.arange(n) + np.random.randint(B) ) % B # allocate equally between all flows mask = tf.one_hot(selected_flows, B, dtype=dtype) mask = mask[:, None, :, None] # n x 1 x B x 1 sample = tf.reduce_sum(sample * mask, -2) # n x N x K return sample, mask
def reverse(self, inputs, **kwargs): """Reverse pass returning the inverse autoregressive transformation.""" if not self.built: self._maybe_build(inputs) net = self.layer(inputs, **kwargs) if net.shape[-1] == self.vocab_size: loc = net scaled_inputs = inputs else: raise ValueError( 'Output of layer does not have compatible dimensions.') loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) outputs = utils.one_hot_add(loc, scaled_inputs) return outputs
def call(self, inputs, **kwargs): """Forward pass for bipartite generation.""" inputs = tf.convert_to_tensor(inputs) batch_ndims = inputs.shape.ndims - 2 mask = tf.reshape(tf.cast(self.mask, inputs.dtype), [1] * batch_ndims + [-1, 1]) masked_inputs = mask * inputs net = self.layer(masked_inputs, **kwargs) if net.shape[-1] == self.vocab_size: loc = net loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) masked_outputs = (1. - mask) * utils.one_hot_minus(inputs, loc) else: raise ValueError( 'Output of layer does not have compatible dimensions.') outputs = masked_inputs + masked_outputs return outputs
def reverse(self, inputs, **kwargs): """Reverse pass for the inverse bipartite transformation.""" if not self.built: self._maybe_build(inputs) inputs = tf.convert_to_tensor(inputs) batch_ndims = inputs.shape.ndims - 2 mask = tf.reshape(tf.cast(self.mask, inputs.dtype), [1] * batch_ndims + [-1, 1]) masked_inputs = mask * inputs net = self.layer(masked_inputs, **kwargs) if net.shape[-1] == self.vocab_size: loc = net scaled_inputs = inputs else: raise ValueError( 'Output of layer does not have compatible dimensions.') loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) masked_outputs = (1. - mask) * utils.one_hot_add(loc, scaled_inputs) outputs = masked_inputs + masked_outputs return outputs
def _initial_call(self, new_inputs, length, **kwargs): """Returns Tensor of shape [..., 1, vocab_size]. Args: new_inputs: Tensor of shape [..., vocab_size], the new input to generate its output. length: Length of final desired sequence. **kwargs: Optional keyword arguments to layer. """ inputs = new_inputs[..., tf.newaxis, :] batch_ndims = inputs.shape.ndims - 2 padded_inputs = tf.pad(inputs, paddings=[[0, 0]] * batch_ndims + [[0, length - 1], [0, 0]]) net = self.layer(padded_inputs, **kwargs) if net.shape[-1] == self.vocab_size: loc = net loc = loc[..., 0:1, :] loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) outputs = utils.one_hot_minus(inputs, loc) else: raise ValueError( 'Output of layer does not have compatible dimensions.') return outputs
def _per_timestep_call(self, current_outputs, new_inputs, length, timestep, **kwargs): inputs = tf.concat([current_outputs, new_inputs[..., tf.newaxis, :]], axis=-2) batch_ndims = inputs.shape.ndims - 2 padded_inputs = tf.pad(inputs, paddings=[[0, 0]] * batch_ndims + [[0, length - timestep - 1], [0, 0]]) net = self.layer(padded_inputs, **kwargs) if net.shape[-1] == self.subset_K: #! loc = net loc = loc[..., :(timestep + 1), :] loc = tf.cast(utils.one_hot_argmax(loc, self.temperature), inputs.dtype) # operate on subset: x = inputs x = tf.gather(x, self.shuffling, axis=-1) x1 = x[..., :self.subset_K] #x1 = super().call(x1) x1 = utils.one_hot_minus( x1, loc) #new_outputs = utils.one_hot_minus(inputs, loc) x2 = x[..., self.subset_K:] x = tf.concat([x1, x2], axis=-1) x = tf.gather(x, self.inverted_shuffling, axis=-1) new_outputs = x else: raise ValueError( 'Output of layer does not have compatible dimensions.') outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]], axis=-2) if not tf.executing_eagerly(): outputs.set_shape([None] * batch_ndims + [timestep + 1, self.vocab_size]) return outputs
def call(self, x): shift = utils.one_hot_argmax(self.logits, self.temperature) shifted_inputs = utils.one_hot_minus(x, shift) scale = utils.one_hot_argmax(self.logits_scale, self.temperature) inverse_scale = utils.multiplicative_inverse(scale, self.K) return utils.one_hot_multiply(shifted_inputs, inverse_scale)
def reverse_static(x, logits, temperature): shift = utils.one_hot_argmax(logits, temperature) return utils.one_hot_add(x, shift)
def call_static(x, logits, temperature): shift = utils.one_hot_argmax(logits, temperature) return utils.one_hot_minus(x, shift)
def reverse(self, x): scale = utils.one_hot_argmax(self.logits_scale, self.temperature) scaled_inputs = utils.one_hot_multiply(x, scale) shift = utils.one_hot_argmax(self.logits, self.temperature) return utils.one_hot_add(shift, scaled_inputs)