Custome GRU with new input, modle can not be built

[Tony-sy]

Hello bodies,
I'm trying to build a new GRU layer that takes the input and a code of zero-one sequenses, the code has the same size of the input, every digit of the code encodes one input value as follows: If the digit is one then the cell will consider the input value to do the callculations, while if the digit is 0, then the cell will consider part of the hidden state h(t-1) to do the callculations. The new cell should be used in API model
I defined the cell and the layer classes. as follows:
class MyGRUCell(Layer):
"""This Cell are ale to deal with missing values in the input
1- adding another input matrex is not possible, GRU needs only one input matrix
2- include the mask of missing values with the values in the same input.
The mask of the missing values looks like extra samples. It has the same shape of the
input.
so the input shape should be
"""
def init(self, units,
Input_activation = 'relu',
activation = 'tanh',
recurrent_activation='hard_sigmoid',
use_bias=True,
kernel_initializer='glorot_uniform',
recurrent_initializer='orthogonal',
bias_initializer='zeros',
kernel_regularizer=None,
recurrent_regularizer=None,
bias_regularizer=None,
kernel_constraint=None,
recurrent_constraint=None,
bias_constraint=None,
dropout=0.,
recurrent_dropout=0.,
implementation=0,
**kwargs):
super(MyGRUCell, self).init(**kwargs)
self.units = units
self.Input_activation = activations.get(Input_activation)
self.activation = activations.get(activation)
self.recurrent_activation = activations.get(recurrent_activation)
self.use_bias = use_bias

self.kernel_initializer = initializers.get(kernel_initializer)
self.recurrent_initializer = initializers.get(recurrent_initializer)
self.bias_initializer = initializers.get(bias_initializer)

self.kernel_regularizer = regularizers.get(kernel_regularizer)
self.recurrent_regularizer = regularizers.get(recurrent_regularizer)
self.bias_regularizer = regularizers.get(bias_regularizer)

self.kernel_constraint = constraints.get(kernel_constraint)
self.recurrent_constraint = constraints.get(recurrent_constraint)
self.bias_constraint = constraints.get(bias_constraint)

self.dropout = min(1., max(0., dropout))
self.recurrent_dropout = min(1., max(0., recurrent_dropout))
self.implementation = implementation
self.state_size = self.units
self._dropout_mask = None
self._recurrent_dropout_mask = None

def build(self, input_shape):
input_dim = input_shape[-1]
print("Building: in_d = ", input_dim)
# add weights (W) to be multiplied by x
self.kernel = self.add_weight(shape=(input_dim, self.units * 3),
name='kernel',
initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer,
constraint=self.kernel_constraint)
# add weights (Uht-1) to e multiplied by
self.recurrent_kernel = self.add_weight(
shape=(self.units, self.units * 4),
name='recurrent_kernel',
initializer=self.recurrent_initializer,
regularizer=self.recurrent_regularizer,
constraint=self.recurrent_constraint)

if self.use_bias:
    self.bias = self.add_weight(shape=(self.units * 4,),
                                name='bias',
                                initializer=self.bias_initializer,
                                regularizer=self.bias_regularizer,
                                constraint=self.bias_constraint)
else:
    self.bias = None

self.recurrent_kernel_x = self.recurrent_kernel[:, :self.units]

self.recurrent_kernel_z = self.recurrent_kernel[:, self.units: self.units * 2]

self.recurrent_kernel_r = self.recurrent_kernel[:, self.units: self.units * 3]

self.recurrent_kernel_h = self.recurrent_kernel[:, self.units * 3:]

if self.use_bias:
    self.bias_x = self.bias[:self.units]
    self.bias_z = self.bias[self.units: self.units * 2]
    self.bias_r = self.bias[self.units: self.units * 3]
    self.bias_h = self.bias[self.units * 3:]
else:
    self.bias_x = None
    self.bias_z = None
    self.bias_r = None
    self.bias_h = None
self.built = True

print("Kernal size =", self.kernel.shape)
print("Built Completed")

def call(self, inputs, states, training=None):
print("Call of the cell . . .")
h_tm1 = states[0] # previous memory
print("Cell input shape: ",inputs.shape)
missing = inputs[:, int(inputs.shape[1]/2) :,:]
Values = inputs[:, : int(inputs.shape[1]/2),:]
print("Values and missing code is separated")
if self.implementation == 1:
## My input
h_as_x = K.dot(h_tm1, self.recurrent_kernel_x)
if self.use_bias:
# wh(t-1)+b
h_as_x = K.bias_add(h_as_x, self.bias_x)

    h_as_input = self.Input_activation(h_as_x)
   
    
    inputs_z = missing * Values + (1-missing) * h_as_input
    inputs_r = missing * Values + (1-missing) * h_as_input
    inputs_h = missing * Values + (1-missing) * h_as_input
    
    x_z = K.dot(inputs_z, self.kernel[:, :self.units])
    x_r = K.dot(inputs_r, self.kernel[:, self.units: self.units * 2])
    x_h = K.dot(inputs_h, self.kernel[:, self.units * 2:])
    if self.use_bias:
        x_z = K.bias_add(x_z, self.bias_z)
        x_r = K.bias_add(x_r, self.bias_r)
        x_h = K.bias_add(x_h, self.bias_h)
        
    h_tm1_z = h_tm1
    h_tm1_r = h_tm1
    h_tm1_h = h_tm1
    
    z = self.recurrent_activation(x_z + K.dot(h_tm1_z, self.recurrent_kernel_z))
    r = self.recurrent_activation(x_r + K.dot(h_tm1_r, self.recurrent_kernel_r))

    hh = self.activation(x_h + K.dot(r * h_tm1_h, self.recurrent_kernel_h))
else:
    # Custom implementation of inputs which are already embedding parameters
    print("Else...")
    matrix_x = K.dot(Values, self.kernel)
    if self.use_bias:
        matrix_x = K.bias_add(matrix_x, self.bias)
    
    matrix_inner = K.dot(h_tm1, self.recurrent_kernel[:, :2 * self.units])

    x_z = matrix_x[:, :self.units]
    x_r = matrix_x[:, self.units: 2 * self.units]
    recurrent_z = matrix_inner[:, :self.units]
    recurrent_r = matrix_inner[:, self.units: 2 * self.units]

    z = self.recurrent_activation(x_z + recurrent_z)
    r = self.recurrent_activation(x_r + recurrent_r)

    x_h = matrix_x[:, 2 * self.units:]
    recurrent_h = K.dot(r * h_tm1,
                        self.recurrent_kernel[:, 2 * self.units:])
    hh = self.activation(x_h + recurrent_h)
h = z * h_tm1 + (1 - z) * hh
print(h.shape)
np.concatenate((h, missing), axis =0)
if training is None:
        h._uses_learning_phase = True
print("Call of the Cell is built")
return h, [h],

def get_config(self):
config = {'units': self.units,
'activation': activations.serialize(self.activation),
'recurrent_activation': activations.serialize(self.recurrent_activation),
'use_bias': self.use_bias,
'kernel_initializer': initializers.serialize(self.kernel_initializer),
'recurrent_initializer': initializers.serialize(self.recurrent_initializer),
'bias_initializer': initializers.serialize(self.bias_initializer),
'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
'recurrent_regularizer': regularizers.serialize(self.recurrent_regularizer),
'bias_regularizer': regularizers.serialize(self.bias_regularizer),
'kernel_constraint': constraints.serialize(self.kernel_constraint),
'recurrent_constraint': constraints.serialize(self.recurrent_constraint),
'bias_constraint': constraints.serialize(self.bias_constraint),
'recurrent_dropout': self.recurrent_dropout,
'implementation': self.implementation}
base_config = super(MyGRUCell, self).get_config()
return dict(list(base_config.items()) + list(config.items()))

print("Cell is Built")
##################
######## Build the class of the GRU
class CMyGRU(RNN):
#@interfaces.legacy_recurrent_support
def init(self, units,
Input_activation= 'relu',
activation='tanh',
recurrent_activation='hard_sigmoid',
use_bias=True,
kernel_initializer='glorot_uniform',
recurrent_initializer='orthogonal',
bias_initializer='zeros',
kernel_regularizer=None,
recurrent_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
recurrent_constraint=None,
bias_constraint=None,
dropout=0.,
recurrent_dropout=0.,
implementation=1,
return_sequences=False,
return_state=False,
go_backwards=False,
stateful=False,
unroll=False,
**kwargs):
if implementation == 0:
warnings.warn('implementation=0 has been deprecated, '
'and now defaults to implementation=1.'
'Please update your layer call.')

    cell = MyGRUCell(units,
               Input_activation=Input_activation,
               activation=activation,
               recurrent_activation=recurrent_activation,
               use_bias=use_bias,
               kernel_initializer=kernel_initializer,
               recurrent_initializer=recurrent_initializer,
               bias_initializer=bias_initializer,
               kernel_regularizer=kernel_regularizer,
               recurrent_regularizer=recurrent_regularizer,
               bias_regularizer=bias_regularizer,
               kernel_constraint=kernel_constraint,
               recurrent_constraint=recurrent_constraint,
               bias_constraint=bias_constraint,
               dropout=dropout,
               recurrent_dropout=recurrent_dropout,
               implementation=implementation)
    super(CMyGRU, self).__init__(cell,
                          return_sequences=return_sequences,
                          return_state=return_state,
                          go_backwards=go_backwards,
                          stateful=stateful,
                          unroll=unroll,
                          **kwargs)
    self.activity_regularizer = regularizers.get(activity_regularizer)
    self.input_spec = [InputSpec(ndim=3)]

def call(self, inputs, mask=None, training=None, initial_state=None):
    print("Call of the Class...")
    print("Call Class input shape:", inputs.shape)
    self.cell._dropout_mask = None
    self.cell._recurrent_dropout_mask = None
    return super(CMyGRU, self).call(inputs,  
                                 mask=mask,
                                 training=training,
                                 initial_state=initial_state)

print("Layer is Built")
#######################

I'm trying to build a neural network model with the new GRU layer as follows:
Input_Features_size = 10
Enter = Input((Input_Features_size*2,1))
print(Enter.shape)
L1 = CMyGRU(units=20, return_sequences = True)(Enter)
L2 = CMyGRU(20, return_sequences = True)(L1, miss= mis_values)
L3 = CMyGRU(10, return_sequences = True )(L2, mis_values)
L4 = CMyGRU(1, return_sequences = False )(L3, mis_values)

Mdl = Model(inputs = Enter, outputs = L4)

opt= optimizers.Adam(lr = 0.001)
Mdl.compile(optimizer = opt, loss = losses.MeanSquaredError(), metrics = ['accuracy'])
Mdl.summary()

unfortunatilly, it gives the following error
ValueError: in user code:

<ipython-input-3-44cc98e553c7>:65 call  *
    mask=mask,
<ipython-input-2-c78e99a9f3df>:106 call  *
    missing = inputs[:, int(inputs.shape[1]/2) :,:]
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\util\dispatch.py:201 wrapper
    return target(*args, **kwargs)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\ops\array_ops.py:1024 _slice_helper
    name=name)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\util\dispatch.py:201 wrapper
    return target(*args, **kwargs)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\ops\array_ops.py:1196 strided_slice
    shrink_axis_mask=shrink_axis_mask)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\ops\gen_array_ops.py:10351 strided_slice
    shrink_axis_mask=shrink_axis_mask, name=name)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\framework\op_def_library.py:744 _apply_op_helper
    attrs=attr_protos, op_def=op_def)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\framework\func_graph.py:593 _create_op_internal
    compute_device)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\framework\ops.py:3485 _create_op_internal
    op_def=op_def)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\framework\ops.py:1975 __init__
    control_input_ops, op_def)
D:\Users\Me\Anaconda3\envs\Work_with_CPU\lib\site-packages\tensorflow\python\framework\ops.py:1815 _create_c_op
    raise ValueError(str(e))

ValueError: Index out of range using input dim 2; input has only 2 dims for '{{node c_my_gru/my_gru_cell/strided_slice}} = StridedSlice[Index=DT_INT32, T=DT_FLOAT, begin_mask=5, ellipsis_mask=0, end_mask=7, new_axis_mask=0, shrink_axis_mask=0](c_my_gru/strided_slice_2, c_my_gru/my_gru_cell/strided_slice/stack, c_my_gru/my_gru_cell/strided_slice/stack_1, c_my_gru/my_gru_cell/strided_slice/stack_2)' with input shapes: [?,1], [3], [3], [3] and with computed input tensors: input[3] = <1 1 1>.

Anyone can find the reason, I think it is related to shape of the input and how it is passed to the inner cell, but i don't know how to fix it.