osl_dynamics.inference.layers

Custom Tensorflow layers.

Module Contents

Classes

DummyLayer	Returns the inputs without modification.
AddRegularizationLossLayer	Adds a regularization loss.
ConcatenateLayer	Wrapper for tf.concat.
MatMulLayer	Wrapper for tf.matmul.
TFRangeLayer	Wrapper for tf.range.
ZeroLayer	Wrapper for tf.zeros.
InverseCholeskyLayer	Layer for getting Cholesky vectors from postive definite symmetric matrices.
BatchSizeLayer	Layer for getting the batch size.
SampleGammaDistributionLayer	Layer for sampling from a Gamma distribution.
SampleNormalDistributionLayer	Layer for sampling from a Normal distribution.
SampleGumbelSoftmaxDistributionLayer	Layer for sampling from a Gumbel-Softmax distribution.
SampleOneHotCategoricalDistributionLayer	Layer for sampling from a Categorical distribution.
SoftmaxLayer	Layer for applying a softmax activation function.
LearnableTensorLayer	Layer to learn a tensor.
VectorsLayer	Layer to learn a set of vectors.
CovarianceMatricesLayer	Layer to learn a set of covariance matrices.
CorrelationMatricesLayer	Layer to learn a set of correlation matrices.
DiagonalMatricesLayer	Layer to learn a set of diagonal matrices.
MatrixLayer	Layer to learn a matrix.
MixVectorsLayer	Mix a set of vectors.
MixMatricesLayer	Layer to mix matrices.
ConcatVectorsMatricesLayer	Layer to concatenate vectors and matrices.
LogLikelihoodLossLayer	Layer to calculate the negative log-likelihood.
AdversarialLogLikelihoodLossLayer	Layer to calculate the negative log-likelihood for SAGE/MAGE.
KLDivergenceLayer	Layer to calculate a KL divergence between two Normal distributions.
KLLossLayer	Layer to calculate the KL loss.
InferenceRNNLayer	RNN inference network.
ModelRNNLayer	RNN generative model.
CategoricalKLDivergenceLayer	Layer to calculate a KL divergence between two categorical distributions.
CategoricalLogLikelihoodLossLayer	Layer to calculate the log-likelihood loss assuming a categorical model.
CategoricalPoissonLogLikelihoodLossLayer	Layer to calculate the log-likelihood loss assuming a categorical model
ConcatEmbeddingsLayer	Layer for getting the concatenated embeddings.
SessionParamLayer	Layer for getting the array specific parameters.
MixSessionSpecificParametersLayer	Class for mixing means and covariances.
GammaExponentialKLDivergenceLayer	Layer to calculate KL divergence between Gamma posterior and exponential
MultiLayerPerceptronLayer	Multi-Layer Perceptron layer.
StaticLossScalingFactorLayer	Layer for calculating the scaling factor for static losses.
HiddenMarkovStateInferenceLayer	Hidden Markov state inference layer.
SeparateLogLikelihoodLayer	Layer to calculate the log-likelihood for different HMM states.
SumLogLikelihoodLossLayer	Layer for summing log-likelihoods.
TFGatherLayer	Wrapper for tf.gather.
AddLayer	Wrapper for tf.add.
EmbeddingLayer	Layer for embeddings.
ShiftForForecastingLayer	Clip two tensors to ensure they align for causal forecasting.
TFConstantLayer	Wrapper for tf.constant.

Functions

add_epsilon(A, epsilon[, diag])	Adds epsilon the the diagonal of batches of square matrices
NormalizationLayer(norm_type, *args, **kwargs)	Returns a normalization layer.
RNNLayer(rnn_type, *args, **kwargs)	Returns a Recurrent Neural Network (RNN) layer.

Attributes

tfb

osl_dynamics.inference.layers.tfb

osl_dynamics.inference.layers.add_epsilon(A, epsilon, diag=False)

Adds epsilon the the diagonal of batches of square matrices or all elements of matrices.

Parameters:

• A (tf.Tensor) – Batches of square matrices or vectors. Shape is (…, N, N) or (…, N).

• epsilon (float) – Small error added to the diagonal of the matrices or every element of the vectors.

• diag (bool, optional) – Do we want to add epsilon to the diagonal only?

osl_dynamics.inference.layers.NormalizationLayer(norm_type, *args, **kwargs)

Returns a normalization layer.

Parameters:

• norm_type (str) – Type of normalization layer. Either 'layer', 'batch' or None.

• args (arguments) – Arguments to pass to the normalization layer.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the normalization layer.

osl_dynamics.inference.layers.RNNLayer(rnn_type, *args, **kwargs)

Returns a Recurrent Neural Network (RNN) layer.

Parameters:

• rnn_type (str) – Type of RNN. Either 'lstm' or 'gru'.

• args (arguments) – Arguments to pass to the normalization layer.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the normalization layer.

class osl_dynamics.inference.layers.DummyLayer

Bases: tensorflow.keras.layers.Layer

Returns the inputs without modification.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.AddRegularizationLossLayer(reg, strength, **kwargs)

Bases: tensorflow.keras.layers.Layer

Adds a regularization loss.

Can be used as a wrapper for a keras regularizer. Inputs are used to calculate the regularization loss and returned without modification.

Parameters:

• reg (str) – Type of regularization. Passed to tf.keras.regularizers.get.

• strength (float) – Strength of regularization. The regularization is multiplied by the strength before being added to the loss.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.ConcatenateLayer(axis, **kwargs)

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.concat.

Parameters:

• axis (int) – Axis to concatenate along.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.MatMulLayer

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.matmul.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.TFRangeLayer(limit, **kwargs)

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.range.

Parameters:

• limit (int) – Upper limit for range.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

class osl_dynamics.inference.layers.ZeroLayer(shape, **kwargs)

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.zeros.

Parameters:

• shape (tuple) – Shape of the zeros tensor.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.InverseCholeskyLayer(epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for getting Cholesky vectors from postive definite symmetric matrices.

Parameters:

• epsilon (float) – Small error added to the diagonal of the matrices.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

class osl_dynamics.inference.layers.BatchSizeLayer

Bases: tensorflow.keras.layers.Layer

Layer for getting the batch size.

Parameters:

kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.SampleGammaDistributionLayer(epsilon, do_annealing, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for sampling from a Gamma distribution.

This layer is a wrapper for tfp.distributions.Gamma.

Parameters:

• epsilon (float) – Error to add to the shape and rate for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, training=None, **kwargs)

This method accepts the shape and rate and outputs the samples.

class osl_dynamics.inference.layers.SampleNormalDistributionLayer(epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for sampling from a Normal distribution.

This layer is a wrapper for tfp.distributions.Normal.

Parameters:

• epsilon (float) – Error to add to the standard deviations for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, training=None, **kwargs)

This method accepts the mean and the standard deviation and outputs the samples.

class osl_dynamics.inference.layers.SampleGumbelSoftmaxDistributionLayer(temperature, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for sampling from a Gumbel-Softmax distribution.

This layer is a wrapper for tfp.distributions.RelaxedOneHotCategorical.

Parameters:

• temperature (float) – Temperature for the Gumbel-Softmax distribution.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

This method accepts logits and outputs samples.

class osl_dynamics.inference.layers.SampleOneHotCategoricalDistributionLayer

Bases: tensorflow.keras.layers.Layer

Layer for sampling from a Categorical distribution.

This layer is a wrapper for tfp.distributions.OneHotCategorical.

call(inputs, **kwargs)

The method accepts logits and outputs samples.

class osl_dynamics.inference.layers.SoftmaxLayer(initial_temperature, learn_temperature, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for applying a softmax activation function.

Parameters:

• initial_temperature (float) – Temperature parameter.

• learn_temperature (bool) – Should we learn the alpha temperature?

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.LearnableTensorLayer(shape, learn=True, initializer=None, initial_value=None, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a tensor.

Parameters:

• shape (tuple) – Shape of the tensor.

• learn (bool, optional) – Should we learn the tensor?

• initializer (tf.keras.initializers.Initializer, optional) – Initializer for the tensor.

• initial_value (float, optional) – Initial value for the tensor if initializer is not passed.

• regularizer (osl-dynamics regularizer, optional) – Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

add_regularization(tensor, static_loss_scaling_factor)

build(input_shape)

call(inputs, static_loss_scaling_factor=1, training=None, **kwargs)

class osl_dynamics.inference.layers.VectorsLayer(n, m, learn, initial_value, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a set of vectors.

Parameters:

• n (int) – Number of vectors.

• m (int) – Number of elements.

• learn (bool) – Should we learn the vectors?

• initial_value (np.ndarray) – Initial value for the vectors. Shape must be (n, m).

• regularizer (osl-dynamics regularizer, optional) –

  Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.CovarianceMatricesLayer(n, m, learn, initial_value, epsilon, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a set of covariance matrices.

A cholesky factor is learnt and used to calculate a covariance matrix as \(C = LL^T\), where \(L\) is the cholesky factor. The cholesky factor is learnt as a vector of free parameters.

Parameters:

• n (int) – Number of matrices.

• m (int) – Number of rows/columns.

• learn (bool) – Should the matrices be learnable?

• initial_value (np.ndarray) – Initial values for the matrices. Shape must be (n, m, m).

• epsilon (float) – Error added to the diagonal of covariances matrices for numerical stability.

• regularizer (osl-dynamics regularizer, optional) –

  Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.CorrelationMatricesLayer(n, m, learn, initial_value, epsilon, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a set of correlation matrices.

A cholesky factor is learnt as a vector of free parameters and used to calculate a correlation matrix.

Parameters:

• n (int) – Number of matrices.

• m (int) – Number of rows/columns.

• learn (bool) – Should the matrices be learnable?

• initial_value (np.ndarray) – Initial values for the matrices. Shape must be (n, m, m).

• epsilon (float) – Error added to the diagonal of correlation matrices for numerical stability.

• regularizer (osl-dynamics regularizer, optional) –

  Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.DiagonalMatricesLayer(n, m, learn, initial_value, epsilon, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a set of diagonal matrices.

The diagonal is forced to be positive using a softplus transformation.

Parameters:

• n (int) – Number of matrices.

• m (int) – Number of rows/columns.

• learn (bool) – Should the matrices be learnable?

• initial_value (np.ndarray) – Initial values for the matrices. Shape must be (n, m, m) or (n, m).

• epsilon (float) – Error added to the diagonal matrices for numerical stability.

• regularizer (osl-dynamics regularizer, optional) –

  Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.MatrixLayer(m, constraint, learn, initial_value, epsilon, regularizer=None, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to learn a matrix.

Parameters:

• m (int) – Number of rows/columns.

• constraint (str) – Either ‘covariance’ or ‘diagonal’.

• learn (bool) – Should the matrix be trainable?

• initial_value (np.ndarray) – Initial value for the matrix.

• epsilon (float) – Error added to the matrices for numerical stability.

• regularizer (osl-dynamics regularizer, optional) –

  Regularizer for the tensor. Must be from inference.regularizers.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

Note

The inputs passed to this method are not used.

class osl_dynamics.inference.layers.MixVectorsLayer

Bases: tensorflow.keras.layers.Layer

Mix a set of vectors.

The mixture is calculated as \(m_t = \displaystyle\sum_j \alpha_{jt} \mu_j\), where \(\mu_j\) are the vectors and \(\alpha_{jt}\) are mixing coefficients.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.MixMatricesLayer

Bases: tensorflow.keras.layers.Layer

Layer to mix matrices.

The mixture is calculated as \(C_t = \displaystyle\sum_j \alpha_{jt} D_j\), where \(D_j\) are the matrices and \(\alpha_{jt}\) are mixing coefficients.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.ConcatVectorsMatricesLayer

Bases: tensorflow.keras.layers.Layer

Layer to concatenate vectors and matrices.

call(inputs, **kwargs)

This method takes a (…, m, n) tensor and (…, m, n, n) tensor.

class osl_dynamics.inference.layers.LogLikelihoodLossLayer(epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the negative log-likelihood.

We assume a multivariate normal probability density. This layer will add the negative log-likelihood to the loss.

Parameters:

• epsilon (float) – Error added to the covariance matrices for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

The method takes the data, mean vector and covariance matrix.

class osl_dynamics.inference.layers.AdversarialLogLikelihoodLossLayer(n_channels, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the negative log-likelihood for SAGE/MAGE.

This layer will add the negative log-likelihood to the loss.

Parameters:

• n_channels (int) – Number of channels.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(y_true, y_pred)

class osl_dynamics.inference.layers.KLDivergenceLayer(epsilon, clip_start=0, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate a KL divergence between two Normal distributions.

Parameters:

• epsilon (float) – Error added to the standard deviations for numerical stability.

• clip_start (int, optional) – Index to clip the sequences inputted to this layer. Default is no clipping.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.KLLossLayer(do_annealing, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the KL loss.

This layer sums KL divergences if multiple values as passed, applies an annealing factor and adds the value to the loss function.

Parameters:

• do_annealing (bool) – Should we perform KL annealing?

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.InferenceRNNLayer(rnn_type, norm_type, act_type, n_layers, n_units, drop_rate, reg, **kwargs)

Bases: tensorflow.keras.layers.Layer

RNN inference network.

Parameters:

• rnn_type (str) – Either 'lstm' or 'gru'.

• norm_type (str) – Either 'layer', 'batch' or None.

• act_type ('str') – Activation type, e.g. 'relu', 'elu', etc.

• n_layers (int) – Number of layers.

• n_units (int) – Number of units/neurons per layer.

• drop_rate (float) – Dropout rate for the output of each layer.

• reg (str) – Regularization for each layer.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.ModelRNNLayer(rnn_type, norm_type, act_type, n_layers, n_units, drop_rate, reg, **kwargs)

Bases: tensorflow.keras.layers.Layer

RNN generative model.

Parameters:

• rnn_type (str) – Either 'lstm' or 'gru'.

• norm_type (str) – Either 'layer', 'batch' or None.

• act_type ('str') – Activation type, e.g. 'relu', 'elu', etc.

• n_layers (int) – Number of layers.

• n_units (int) – Number of units/neurons per layer.

• drop_rate (float) – Dropout rate for the output of each layer.

• reg (str) – Regularization for each layer.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.CategoricalKLDivergenceLayer(clip_start=0, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate a KL divergence between two categorical distributions.

Parameters:

• clip_start (int, optional) – Index to clip the sequences inputted to this layer. Default is no clipping.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.CategoricalLogLikelihoodLossLayer(n_states, epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the log-likelihood loss assuming a categorical model.

Parameters:

• n_states (int) – Number of states.

• epsilon (float) – Error added to the covariances for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.CategoricalPoissonLogLikelihoodLossLayer(n_states, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the log-likelihood loss assuming a categorical model with Poisson observation model.

Parameters:

• n_states (int) – Number of states.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.ConcatEmbeddingsLayer

Bases: tensorflow.keras.layers.Layer

Layer for getting the concatenated embeddings.

The concatenated embeddings are obtained by concatenating embeddings and spatial embeddings.

call(inputs)

class osl_dynamics.inference.layers.SessionParamLayer(param, epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for getting the array specific parameters.

This layer adds deviations to the group spatial parameters.

Parameters:

• param (str) – Which parameter are we using? Must be 'means' or 'covariances'.

• epsilon (float) – Error added to the diagonal of covariances for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs)

class osl_dynamics.inference.layers.MixSessionSpecificParametersLayer

Bases: tensorflow.keras.layers.Layer

Class for mixing means and covariances.

The mixture is calculated as

• \(m_t = \displaystyle\sum_j \alpha_{jt} \mu_j^{s_t}\)

• \(C_t = \displaystyle\sum_j \alpha_{jt} D_j^{s_t}\)

where \(s_t\) is the array at time \(t\).

call(inputs)

class osl_dynamics.inference.layers.GammaExponentialKLDivergenceLayer(epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate KL divergence between Gamma posterior and exponential prior for deviation magnitude.

Parameters:

• epsilon (float) – Error added to the standard deviations for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, static_loss_scaling_factor=1, **kwargs)

class osl_dynamics.inference.layers.MultiLayerPerceptronLayer(n_layers, n_units, norm_type, act_type, drop_rate, regularizer=None, regularizer_factor=0.0, **kwargs)

Bases: tensorflow.keras.layers.Layer

Multi-Layer Perceptron layer.

Parameters:

• n_layers (int) – Number of layers.

• n_units (int) – Number of units/neurons.

• norm_type (str) – Normalization layer type. Can be 'layer', 'batch' or None.

• act_type (str) – Activation type.

• drop_rate (float) – Dropout rate.

• regularizer (str, optional) – Regularizer type.

• regularizer_factor (float, optional) – Regularizer factor.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the base class.

call(inputs, static_loss_scaling_factor=1, training=None, **kwargs)

class osl_dynamics.inference.layers.StaticLossScalingFactorLayer(**kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for calculating the scaling factor for static losses.

When calculating loss, we sum over the sequence length (time dimension) and average over the sequences. If we add a static quantity to each time point we need to rescale it to account for the summation over time. The scaling factor is given by

\[\text{static_loss_scaling_factor} = \frac{1}{\text{batch_size} \times \text{n_batches}}\]

call(inputs, **kwargs)

class osl_dynamics.inference.layers.HiddenMarkovStateInferenceLayer(n_states, initial_trans_prob, learn, use_stationary_distribution=False, **kwargs)

Bases: tensorflow.keras.layers.Layer

Hidden Markov state inference layer.

This layer uses the Baum-Welch algorithm to calculate the posterior for the hidden state in a Hidden Markov Model (HMM).

Parameters:

• n_states (int) – Number of states.

• initial_trans_prob (np.ndarray) – Initial transition probability matrix. Shape must be (n_states, n_states.)

• learn (bool) – Should we learn the transition probability matrix?

• use_stationary_distribution (bool, optional) – Should we use the stationary distribution (estimated from the transition probability matrix) for the initial state probabilities?

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the normalization layer.

get_stationary_distribution()

get_trans_prob()

_baum_welch(log_B)

_trans_prob_update(log_gamma, log_xi)

call(log_B, **kwargs)

class osl_dynamics.inference.layers.SeparateLogLikelihoodLayer(n_states, epsilon, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer to calculate the log-likelihood for different HMM states.

Parameters:

• n_states (int) – Number of states.

• epsilon (float) – Error added to the covariance matrices for numerical stability.

• kwargs (keyword arguments, optional) – Keyword arguments to pass to the normalization layer.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.SumLogLikelihoodLossLayer

Bases: tensorflow.keras.layers.Layer

Layer for summing log-likelihoods.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.TFGatherLayer(axis, batch_dims=0, **kwargs)

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.gather.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.AddLayer

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.add.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.EmbeddingLayer(input_dim, output_dim, unit_norm, **kwargs)

Bases: tensorflow.keras.layers.Layer

Layer for embeddings.

Parameters:

• input_dim (int) – Input dimension.

• output_dim (int) – Output dimension.

• unit_norm (bool, optional) – Should the embeddings be unit norm?

property embeddings

call(inputs, **kwargs)

class osl_dynamics.inference.layers.ShiftForForecastingLayer(clip, **kwargs)

Bases: tensorflow.keras.layers.Layer

Clip two tensors to ensure they align for causal forecasting.

Parameters:

clip (int) – Number of elements to clip.

call(inputs, **kwargs)

class osl_dynamics.inference.layers.TFConstantLayer(values, **kwargs)

Bases: tensorflow.keras.layers.Layer

Wrapper for tf.constant.

call(inputs, **kwargs)