code for DCT transform

import keras.backend as K
import tensorflow as tf

def samplewise_tf_dct(x_sample):
    print('=======> input of "samplewise_tf_dct"')
    print(type(x_sample))
    print((x_sample.shape))
    print(x_sample)
    
    output_list = [] 
    for idx_channel in range(x_sample.shape[-1]):
        output_list.append(tf.spectral.dct(K.transpose(tf.spectral.dct(x_sample[:, :, idx_channel])))) # a 2D DCT    
    
    print('=======> output of "samplewise_tf_dct"')
    print(K.stack(output_list, axis=-1))
    print('<==========================================================>')
    
    return K.stack(output_list, axis=-1)

def dct_layer_function(x_batch):
    print('=======> input of "dct_layer_function"')
    print(type(x_batch))
    print((x_batch.shape))
    print(x_batch)
    
    print('=======> output of "dct_layer_function"')
    print(K.map_fn(samplewise_tf_dct, x_batch))
    print('<==========================================================>')
    
    return K.map_fn(samplewise_tf_dct, x_batch)

####################################################################
# between two blocks of code in jupyter notebook
####################################################################

from keras.models import Model
from keras import layers
from keras.layers import Lambda

input_of_net = layers.Input(shape=(27, 27, 3), name='input_of_net')

x = layers.Conv2D(32, (3, 3), strides=(2, 2), kernel_initializer='glorot_normal', name='block1_conv1')(input_of_net)
x = layers.BatchNormalization(name='block1_conv1_bn')(x)
x = layers.Activation('relu', name='block1_conv1_act')(x)

dct_layer = Lambda(function = dct_layer_function)
x = dct_layer(x)

x = layers.Conv2D(64, (3, 3), kernel_initializer='glorot_normal', name='block1_conv2')(x)
x = layers.BatchNormalization(name='block1_conv2_bn')(x)
x = layers.Activation('relu', name='block1_conv2_act')(x)

x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(2, activation = 'sigmoid')(x)
model = Model(inputs = input_of_net, outputs = x)

model.summary()

####################################################################
# between two blocks of code in jupyter notebook
####################################################################

Result:
=======> input of "dct_layer_function"
<class 'tensorflow.python.framework.ops.Tensor'>
(?, 13, 13, 32)
Tensor("block1_conv1_act_10/Relu:0", shape=(?, 13, 13, 32), dtype=float32)
=======> output of "dct_layer_function"
=======> input of "samplewise_tf_dct"
<class 'tensorflow.python.framework.ops.Tensor'>
(13, 13, 32)
Tensor("lambda_11/map/while/TensorArrayReadV3:0", shape=(13, 13, 32), dtype=float32)
=======> output of "samplewise_tf_dct"
Tensor("lambda_11/map/while/stack:0", shape=(13, 13, 32), dtype=float32)
<==========================================================>
Tensor("lambda_11/map/TensorArrayStack/TensorArrayGatherV3:0", shape=(?, 13, 13, 32), dtype=float32)
<==========================================================>
=======> input of "samplewise_tf_dct"
<class 'tensorflow.python.framework.ops.Tensor'>
(13, 13, 32)
Tensor("lambda_11/map_1/while/TensorArrayReadV3:0", shape=(13, 13, 32), dtype=float32)
=======> output of "samplewise_tf_dct"
Tensor("lambda_11/map_1/while/stack:0", shape=(13, 13, 32), dtype=float32)
<==========================================================>
=======> input of "dct_layer_function"
<class 'tensorflow.python.framework.ops.Tensor'>
(?, 13, 13, 32)
Tensor("lambda_11/Placeholder:0", shape=(?, 13, 13, 32), dtype=float32)
=======> output of "dct_layer_function"
=======> input of "samplewise_tf_dct"
<class 'tensorflow.python.framework.ops.Tensor'>
(13, 13, 32)
Tensor("lambda_11/map_2/while/TensorArrayReadV3:0", shape=(13, 13, 32), dtype=float32)
=======> output of "samplewise_tf_dct"
Tensor("lambda_11/map_2/while/stack:0", shape=(13, 13, 32), dtype=float32)
<==========================================================>
Tensor("lambda_11/map_2/TensorArrayStack/TensorArrayGatherV3:0", shape=(?, 13, 13, 32), dtype=float32)
<==========================================================>
=======> input of "samplewise_tf_dct"
<class 'tensorflow.python.framework.ops.Tensor'>
(13, 13, 32)
Tensor("lambda_11/map_3/while/TensorArrayReadV3:0", shape=(13, 13, 32), dtype=float32)
=======> output of "samplewise_tf_dct"
Tensor("lambda_11/map_3/while/stack:0", shape=(13, 13, 32), dtype=float32)
<==========================================================>
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_of_net (InputLayer)    (None, 27, 27, 3)         0         
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 13, 13, 32)        896       
_________________________________________________________________
block1_conv1_bn (BatchNormal (None, 13, 13, 32)        128       
_________________________________________________________________
block1_conv1_act (Activation (None, 13, 13, 32)        0         
_________________________________________________________________
lambda_11 (Lambda)           (None, 13, 13, 32)        0         
_________________________________________________________________
block1_conv2 (Conv2D)        (None, 11, 11, 64)        18496     
_________________________________________________________________
block1_conv2_bn (BatchNormal (None, 11, 11, 64)        256       
_________________________________________________________________
block1_conv2_act (Activation (None, 11, 11, 64)        0         
_________________________________________________________________
global_average_pooling2d_3 ( (None, 64)                0         
_________________________________________________________________
dense_3 (Dense)              (None, 2)                 130       
=================================================================
Total params: 19,906
Trainable params: 19,714
Non-trainable params: 192
_________________________________________________________________