Optionally, you an also implement the method get_config and the class method from_config in order to support serialization -- just like with any Keras object. Note that we don't have to implement from_config in the example above since the constructor arguments of the class the keys in the config returned by get_config are the same. In this case, the default from_config works fine.