Add a custom layer with a bitwise operation

[AnouarITI]

Hello, I have a use case where I want to introduce a bitwise operation or as a custom layer. As a 1st step, I trained a simple model with 8-bit integers as follows:

from qkeras import *

qmodel = Sequential()
qmodel.add(QDense(512, 
                  kernel_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)", 
                  bias_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)",
                  input_shape=(784,)))

qmodel.add(QActivation("quantized_relu(bits=8, integer=7)"))

qmodel.add(QDense(256,  
                  kernel_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)", 
                  bias_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)"))

qmodel.add(QActivation("quantized_relu(bits=8, integer=7)"))

qmodel.add(QDense(10, 
                  kernel_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)", 
                  bias_quantizer="quantized_bits(bits=8, integer=7, keep_negative=True, alpha=1)"))

qmodel.add(Activation("softmax"))
qmodel.summary()

Then I want to add another layer that performs the following operation (x | 5) % 256, which looks similar to this in Keras:

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(CustomLayer, self).__init__(**kwargs)

    def call(self, x):
        result = tf.math.mod((x | 5), 256)
        return result 

How can I proceed if I want to make the same in QKeras to be able to perform QAT with the custom layer?