Layers

VAE.layers

Collection of Keras layers.

VAE.layers.Attention

Attention(activation='softmax', permute=None, **kwargs)

Bases: Layer

Dot-product attention layer.

This layers builds on the attention layer :class:keras.layers.Attention. The matrix multiplication in the attention is applied to the two inner dimensions and broadcasted otherwise. Use permute to reorder the input before attention. The output has again the same order of dimenions as the input.

Parameters:

• activation (str, default: 'softmax' ) –

  Name of activation function applied to the score. Defaults to 'softmax'.

• permute (list[int], default: None ) –

  Permutation applied to the dimensions of the input tensors before attention. The output has the same order of dimenions as the input. Defaults to None, meaning that no permutation is applied.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

Source code in VAE/layers.py

def __init__(self, activation: str = 'softmax', permute: list[int] = None, **kwargs):
    super().__init__(**kwargs)

    self.activation = ks.activations.get(activation)
    self.permute = permute
    self.ipermute = None

VAE.layers.Attention.call

call(inputs)

Call the layer.

Parameters:

• inputs (tuple[Tensor, Tensor, Tensor]) –

  Tuple of three tensors representing the query, key, and value, respectively. The shapes of the three tensors is (batch_size, set_size, time_length, channels).

Returns:

• Tensor –

  Tensor of shape (batch_size, set_size, time_length, channels).

Source code in VAE/layers.py

def call(self, inputs: tuple[tf.Tensor, tf.Tensor, tf.Tensor]) -> tf.Tensor:
    """Call the layer.

    Parameters:
        inputs:
            Tuple of three tensors representing the query, key, and value, respectively. The shapes of the three
            tensors is `(batch_size, set_size, time_length, channels)`.

    Returns:
        Tensor of shape `(batch_size, set_size, time_length, channels)`.

    """
    q, k, v = inputs

    if self.permute is not None:
        q = tf.transpose(q, perm=self.permute)
        k = tf.transpose(k, perm=self.permute)
        v = tf.transpose(v, perm=self.permute)

    qk = tf.matmul(q, k, transpose_b=True)
    qk = self.activation(qk)
    qkv = tf.matmul(qk, v)

    if self.ipermute is not None:
        qkv = tf.transpose(qkv, perm=self.ipermute)

    return qkv

VAE.layers.AttentionMasked

AttentionMasked(activation='softmax', masked=True, permute=None, **kwargs)

Bases: Layer

Dot-product attention layer.

This layers builds on the attention layer :class:keras.layers.Attention. Each input tensor of shape (batch_size, set_size, time_length, channels) is first reshaped into a tensor of shape (batch_size, set_size * time_length, channels). The dot-product attention is applied on the last dimension. If masked=True, score values with a query time index larger than the key time index are masked out.

Parameters:

• activation (str, default: 'softmax' ) –

  Activation function applied to the scores. Defaults to 'softmax'.

• masked (bool, default: True ) –

  Whether a causal masked is applied to the scores. Defaults to True.

• permute (tuple[int, int, int], default: None ) –

  More general permutation of the dimensions of the input tensors before attention. Note that the time index of the causal mask always refers to the second dimension of the permuted dimensions. Defaults to None, meaning that no permutation is applied.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

Source code in VAE/layers.py

def __init__(self,
             activation: str = 'softmax',
             masked: bool = True,
             permute: tuple[int, int, int] = None,
             **kwargs):
    super().__init__(**kwargs)
    self.activation = ks.activations.get(activation)
    self.masked = masked
    self.permute = permute

VAE.layers.AttentionMasked.call

call(inputs)

Apply the the dot-product attention to the inputs.

Parameters:

• inputs (tuple[Tensor, Tensor, Tensor]) –

  Tuple of three tensors representing the query, key, and value, respectively. The shapes of the three tensors is (batch_size, set_size, time_length, channels).

Returns:

• Tensor –

  Tensor of shape (batch_size, set_size, time_length, channels).

Source code in VAE/layers.py

def call(self, inputs: tuple[tf.Tensor, tf.Tensor, tf.Tensor]) -> tf.Tensor:
    """Apply the the dot-product attention to the inputs.

    Parameters:
        inputs:
            Tuple of three tensors representing the query, key, and value, respectively. The shapes of the three
            tensors is `(batch_size, set_size, time_length, channels)`.

    Returns:
        Tensor of shape `(batch_size, set_size, time_length, channels)`.

    """
    q, k, v = inputs

    # permute dimensions
    if self.permute is not None:
        q = tf.transpose(q, perm=self.permute)
        k = tf.transpose(k, perm=self.permute)
        v = tf.transpose(v, perm=self.permute)

    q_shape = tf.shape(q)
    k_shape = tf.shape(k)
    v_shape = tf.shape(v)

    q = tf.reshape(q, shape=(-1, q_shape[1] * q_shape[2], q_shape[3]))
    k = tf.reshape(k, shape=(-1, k_shape[1] * k_shape[2], k_shape[3]))
    v = tf.reshape(v, shape=(-1, v_shape[1] * v_shape[2], v_shape[3]))

    # get dot product of queries and keys
    qk = tf.matmul(q, k, transpose_b=True)
    # scale
    qk /= tf.sqrt(tf.cast(q_shape[-1], dtype=qk.dtype))
    # apply mask to dot product: set masked positions to large negative values results in near zero activation
    if self.mask is not None:
        qk -= 1e9 * self.mask
    # activation
    qk = self.activation(qk)

    # get output of attention
    qkv = tf.matmul(qk, v)

    # reshape output
    qkv = tf.reshape(qkv, shape=v_shape)

    # restore order of dimensions in output
    if self.ipermute is not None:
        qkv = tf.transpose(qkv, perm=self.ipermute)

    return qkv

VAE.layers.Film

Film(activation=None, use_scale=True, use_offset=True, use_bias=True, shape=None, kernel_initializer='glorot_uniform', bias_initializer='zeros', **kwargs)

Bases: Layer

Feature-wise linear modulation.

Feature-wise linear modulation of the last dimension of the input tensor.

A linear layer is applied to the last dimension of the condition tensor and the output gives the scale that is multiplied with the input tensor. The scale can be activated by a non-linearity prior to the multiplication.

Another linear layer is applied to the last dimension of the condition tensor and the output gives the offset that is added to the input tensor.

The rank of the condition tensor can be smaller than the rank of the input tensor. In this case, the scale is applied according to shape parameter.

Note: In this version of the FiLM layer, the shape parameter affects both, the scale and the offset.

Parameters:

• activation (str, default: None ) –

  Activation function in the linear layer of the scale.

• use_scale (bool, default: True ) –

  Whether to use the scale.

• use_offset (bool, default: True ) –

  Whether to use the offset.

• use_bias (bool, default: True ) –

  Whether to use a bias in the linear layers of the scale and offset.

• shape (list[int], default: None ) –

  Shape of the modulation. If None, the modulation is applied to the last dimension of the input. This is the default case from [1]. If shape is a tuple of length n (where n = rank(input) - rank(condition) + 1), the modulation is applied to the last n dimensions of the input according to the specified shape. Setting values to None in the shape tuple results in the modulation being applied to this dimension. Setting values to 1 in the shape results in the modulation being broadcasted to this dimension.Default is None.

• kernel_initializer (str, default: 'glorot_uniform' ) –

  Initializer for the kernel of the linear layers.

• bias_initializer (str, default: 'zeros' ) –

  Initializer for the bias of the linear layers.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

References: [1] Perez et al. (2017): https://arxiv.org/abs/1709.07871.

Source code in VAE/layers.py

def __init__(self,
             activation: str = None,
             use_scale: bool = True,
             use_offset: bool = True,
             use_bias: bool = True,
             shape: list[int] = None,
             kernel_initializer: str = 'glorot_uniform',
             bias_initializer: str = 'zeros',
             **kwargs):
    self.activation = ks.activations.get(activation)
    self.use_scale = use_scale
    self.use_offset = use_offset
    self.use_bias = use_bias
    self.shape = shape
    self.kernel_initializer = ks.initializers.get(kernel_initializer)
    self.bias_initializer = ks.initializers.get(bias_initializer)
    super().__init__(**kwargs)

VAE.layers.Film.call

call(inputs)

Apply FiLM to input tensor.

Parameters:

• inputs (tuple[Tensor, Tensor]) –

  Tuple of two tensors of arbitrary shape. The first tensor is the input. The second tensor is the condition. The k - 1 leading dimensions of the condition tensor must be broadcastable to the first k- 1 leading dimensions of the input tensor, where k is the rank of the condition tensor.

Returns:

• Tensor –

  Tensor of the same shape as input tensor.

Source code in VAE/layers.py

def call(self, inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
    """Apply FiLM to input tensor.

    Parameters:
        inputs:
            Tuple of two tensors of arbitrary shape. The first tensor is the input. The second tensor is the
            condition. The `k - 1` leading dimensions of the condition tensor must be broadcastable to the first
            `k- 1` leading dimensions of the input tensor, where `k` is the rank of the condition tensor.

    Returns:
        Tensor of the same shape as input tensor.

    """
    x, cond = inputs

    # scale input
    if self.use_scale:
        gamma = tf.tensordot(cond, self.gamma_kernel, axes=[[-1], [0]])
        if self.use_bias:
            gamma += self.gamma_bias

        if self.activation:
            gamma = self.activation(gamma)

        x *= gamma

    # add offset
    if self.use_offset:
        beta = tf.tensordot(cond, self.beta_kernel, axes=[[-1], [0]])
        if self.use_bias:
            beta += self.beta_bias

        x += beta

    return x

VAE.layers.GumbelSoftmax

GumbelSoftmax(axis=-1, temperature=1.0, hard=False, noise_shape=None, **kwargs)

Bases: Layer

Random sampling from Gumbel softmax distribution.

This layer is used to sample from a Gumbel softmax distribution.

Parameters:

• axis (int, default: -1 ) –

  The axis along which to apply the Gumbel softmax. Default is last axis.

• temperature (float, default: 1.0 ) –

  The temperature of the Gumbel softmax. Default is 1.

• hard (bool, default: False ) –

  Whether to sample from the hard or soft Gumbel distribution.

• noise_shape (list[int], default: None ) –

  The shape of the random noise. Must be of the same length as the number of dimensions in the input. None values in the tuple can be used to infer the shape from the input shape. If None, the noise shape will be equal to the shape of the input. Default is None.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

Source code in VAE/layers.py

def __init__(self,
             axis: int = -1,
             temperature: float = 1.,
             hard: bool = False,
             noise_shape: list[int] = None,
             **kwargs):
    self.axis = axis
    self.temperature = temperature
    self.hard = hard
    self.noise_shape = noise_shape
    super().__init__(**kwargs)

VAE.layers.GumbelSoftmax.call

call(inputs)

Apply Gumbel softmax to input tensor.

Parameters:

• inputs (Tensor) –

  Tensor of arbitrary shape. The input is expected to be logits.

Returns:

• Tensor –

  Tensor of the same shape as input tensor.

Source code in VAE/layers.py

def call(self, inputs: tf.Tensor) -> tf.Tensor:
    """Apply Gumbel softmax to input tensor.

    Parameters:
        inputs:
            Tensor of arbitrary shape. The input is expected to be logits.

    Returns:
        Tensor of the same shape as input tensor.

    """
    x = inputs
    shape = complete_shape(x, self.noise_shape)
    u = tf.random.uniform(shape, minval=0, maxval=1)
    eps = K.epsilon()
    g = -tf.math.log(-tf.math.log(u + eps) + eps)
    y = x + g
    y = tf.nn.softmax(y / self.temperature, axis=self.axis)
    if self.hard:
        y_hard = tf.one_hot(tf.argmax(y, axis=self.axis), depth=tf.shape(y)[self.axis], axis=self.axis)
        y_hard = tf.cast(y_hard, x.dtype)
        y = tf.stop_gradient(y_hard - y) + y
    return y

VAE.layers.RandomSampling

RandomSampling(noise_shape=None, **kwargs)

Bases: Layer

Random sampling from normal distribution.

This layer samples from an isotropic Gaussian with mean z_mean and log variance z_log_var.

Parameters:

• noise_shape (list[int], default: None ) –

  The shape of the random noise. Must be of the same length as the number of dimensions in the input and be broadcastable to the shape of the input. If None, the noise shape will be equal to the shape of the input. Default is None.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

Source code in VAE/layers.py

def __init__(self, noise_shape: list[int] = None, **kwargs):
    self.noise_shape = noise_shape
    super().__init__(**kwargs)

VAE.layers.RandomSampling.call

call(inputs)

Sample from normal distribution.

Parameters:

• inputs (tuple[Tensor, Tensor]) –

  Tuple of two or three tensors of the same shape. The first tensor is the mean of the normal distribution, the second tensor is the logarithm of the variance of the normal distribution. If a third tensor is given, it is used as the random sample. Otherwise, a random sample is generated.

Returns:

• Tensor –

  Tensor of the same shape as input tensors.

Source code in VAE/layers.py

def call(self, inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
    """Sample from normal distribution.

    Parameters:
        inputs:
            Tuple of two or three tensors of the same shape. The first tensor is the mean of the normal
            distribution, the second tensor is the logarithm of the variance of the normal distribution. If a third
            tensor is given, it is used as the random sample. Otherwise, a random sample is generated.

    Returns:
        Tensor of the same shape as input tensors.

    """
    if len(inputs) == 2:
        z_mean, z_log_var = inputs
        shape = complete_shape(z_mean, self.noise_shape)
        sigma = tf.random.normal(shape=shape)
    else:
        z_mean, z_log_var, sigma = inputs

    return z_mean + tf.exp(0.5 * z_log_var) * sigma

VAE.layers.Split

Split(size_splits, axis=0, **kwargs)

Bases: Layer

Split input tensor into smaller chunks.

Parameters:

• size_splits (list[int]) –

  Containing the sizes of each output tensor along axis.

• axis (int, default: 0 ) –

  The dimension along which to split. Defaults to 0.

• **kwargs –

  Additional keyword arguments passed to the Layer superclass.

Source code in VAE/layers.py

def __init__(self, size_splits: list[int], axis: int = 0, **kwargs):
    self.size_splits = size_splits
    self.axis = axis
    super().__init__(**kwargs)

VAE.layers.Split.call

call(inputs)

Split input tensor into smaller chunks.

Parameters:

• inputs (Tensor) –

  A tensor of arbitrary shape.

Returns:

• list[Tensor] –

  A list of tensors.

Source code in VAE/layers.py

def call(self, inputs: tf.Tensor) -> list[tf.Tensor]:
    """Split input tensor into smaller chunks.

    Parameters:
        inputs:
            A tensor of arbitrary shape.

    Returns:
        A list of tensors.

    """
    return tf.split(inputs, self.size_splits, self.axis)

VAE.layers.example_Film

example_Film()

Example of Film layer.

Source code in VAE/layers.py

def example_Film():
    """Example of Film layer."""
    input_shape = (1, 16, 12)
    cond_shape = (10, )
    x = tf.constant(1., shape=(32, ) + input_shape)
    c = tf.constant(1., shape=(32, ) + cond_shape)
    x_in = ks.layers.Input(shape=input_shape)
    cond_in = ks.layers.Input(shape=cond_shape)
    out = Film(
        use_scale=True,
        use_offset=True,
        use_bias=True,
        shape=(1, None, None),
    )([x_in, cond_in])
    model = ks.Model(inputs=[x_in, cond_in], outputs=out)
    model.summary()
    _ = model.predict([x, c])
    for w in model.weights:
        print(w.name, ':', w.shape)

VAE.layers.example_GumbelSoftmax

example_GumbelSoftmax()

Example of GumbelSoftmax layer.

Source code in VAE/layers.py

def example_GumbelSoftmax():
    """Example of GumbelSoftmax layer."""
    input_shape = (5, 4)
    x = tf.zeros((2, ) + input_shape)
    x_in = ks.layers.Input(shape=input_shape)
    out = GumbelSoftmax(axis=-1, hard=True, noise_shape=(None, 1, None))(x_in)
    model = ks.Model(inputs=x_in, outputs=out)
    y = model.predict(x)
    print(y)

VAE.layers.example_RandomSampling

example_RandomSampling()

Example of RandomSampling layer.

Source code in VAE/layers.py

def example_RandomSampling():
    """Example of RandomSampling layer."""
    input_shape = (5, 4)
    z_mean = tf.zeros((2, ) + input_shape)
    z_log_var = tf.zeros((2, ) + input_shape)
    z_mean_in = ks.layers.Input(shape=input_shape)
    z_log_var_in = ks.layers.Input(shape=input_shape)
    out = RandomSampling(noise_shape=(None, 1, None))([z_mean_in, z_log_var_in])
    model = ks.Model(inputs=[z_mean_in, z_log_var_in], outputs=out)
    z = model.predict([z_mean, z_log_var])
    print(z)