custom_layer.py

from tensorflow import keras
import tensorflow as tf

class ChildDense(keras.layers.Layer):
  
  def __init__(self, units, activation=None):
    super().__init__()
    self.units = units
    self.activation = activation
      
  def build(self, input_shape):
    # Shape of the input data is referred by input_shape.
    input_dim = input_shape[-1]
    
    self.W = self.add_weight(shape=(input_dim, self.units), initializer='random_normal')
    self.b = self.add_weight(shape=(self.units, ), initializer='zeros')
    
  def call(self, inputs):
    y = tf.matmul(inputs, self.W) + self.b
    if self.activation is not None:
      y = self.activation(y)
    return y