class MyCustomLayer(layers.Layer):
    def __init__(self, units=32, activation=None):
        super(MyCustomLayer, self).__init__()
        self.units = units
        self.activation = keras.activations.get(activation)

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                 initializer='zeros',
                                 trainable=True)

    def call(self, inputs):
        return self.activation(tf.matmul(inputs, self.w) + self.b)

inputs = keras.Input(shape=(784,))
x = MyCustomLayer(64, activation='relu')(inputs)
outputs = MyCustomLayer(10, activation='softmax')(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        __
        __