time2vec_sample.py

import tensorflow.keras as keras
import tensorflow.keras.backend as K


class Time2Vec(keras.layers.Layer):
    def __init__(self, kernel_size):
        super().__init__(trainable=True, name='Time2Vec')
        self.k = kernel_size

    def build(self, input_shape):
        # trend
        self.wb = self.add_weight(name='wb', shape=(
            input_shape[1],), initializer='uniform', trainable=True)
        self.bb = self.add_weight(name='bb', shape=(
            input_shape[1],), initializer='uniform', trainable=True)
        # periodic
        self.wa = self.add_weight(name='wa', shape=(
            1, input_shape[1], self.k), initializer='uniform', trainable=True)

        self.ba = self.add_weight(name='ba', shape=(
            1, input_shape[1], self.k), initializer='uniform', trainable=True)
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        bias = self.wb * inputs + self.bb
        dp = K.dot(inputs, self.wa) + self.ba
        wgts = K.sin(dp)

        ret = K.concatenate([K.expand_dims(bias, -1), wgts], -1)
        ret = K.reshape(ret, (-1, inputs.shape[1]*(self.k+1)))
        return ret

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1]*(self.k + 1))