How could I assign a tensor to a slice of another tensor with a differentiable tensorflow method

Question

I need to implement a tensorsketch method in my CNN (https://arxiv.org/pdf/1511.06062.pdf) but I have assignement issues.

I am using tensorflow 2 (with python 3) and I already tried several methods but still arrive to the same issue: ValueError: Tried to convert 'input' to a tensor and failed. Error: None values not supported.

My code is the Following:

def while_condition3(j,E,x,h,s,d):
    return tf.less(j, d)

def body3(j,E,x,h,s,d): 
    z = tf.multiply(s[0,j],x[0,j])
    z = K.reshape(z, shape=(-1, 1))
    g = tf.add(E[:,:,y2],z ,name='ADD')
    g = K.reshape(g, shape=(-1, 1)) 
    #g=tf.convert_to_tensor(g)
    E[:,:,h[0,j]].assign(g) # <--------- 3 = error
    return [tf.add(j, 1),E,x,h,s,d]

def sig3(x,h,s,d):
    E = tf.Variable(tf.zeros([tf.shape(x)[0],d]), tf.float32,validate_shape=False)
    j = tf.constant(0)
    r = tf.while_loop(while_condition3, body3, loop_vars=[j,E,x,h,s,d]) # <--------- 2
    return r[1]

def countSh(i,d,x1b):#shape de xi = (B,1,512) # see here https://arxiv.org/pdf/1511.06062.pdf the algorithm 2
    tt = time.time()
    h1 = tf.get_variable( 'h1_{}'.format(tt), [1,K.int_shape(x1b)[1]], initializer=tf.random_uniform_initializer(0, d), trainable=False)
    h1 = tf.cast(h1, 'int32')
    s1 = tf.get_variable('s1_{}'.format(tt), [1,K.int_shape(x1b)[1]], initializer=tf.random_uniform_initializer(0, 2), trainable=False)
    s1 = tf.cast(tf.floor(s1) * 2 - 1, 'float32') # 1 or -1
    E1 = sig3(x1b,h1,s1,d) # <--------- 1
    return E1

The problem is in body3, the line E[:,:,y2].assign(g). I already tried to create E out of the graph but either it was too long because it needs 2 for loops, or there was an error: TypeError: 'Tensor' object does not support item assignment and 'Tensor' object cannot be interpreted as an integer. And can't use just numpy because I don't know the shape.

I also tried what was proposed in this post: How to do slice assignment in Tensorflow but it didn't work.

And finally, I have seen that tf.assign shouldn't be used because it is not differentiable. So is there something that I can use?

EDIT: Here is an example for execution

def forAll(x1r):
    d=6
    F = []#K.zeros((K.shape(x1r)[0],K.shape(x1r)[1],d))
    for i in range(K.int_shape(x1r)[1]): 
        print('pixel ',i) #in 36 because for each pixel
        T = countSh(i,d,x1r[:,i,:])
        F.append(T)
    F=tf.stack(F) #non differentiable
    return F

def compact(modelCN,d,interpolSize,noClasse):
    input = Input(shape=(interpolSize, interpolSize, 3), name='input')
    x1 = modelCN(input) # output cnn1 ##shape = (?,6,6,512)
    b=K.int_shape(x1)[1]*K.int_shape(x1)[2]
    x1r = keras.layers.Reshape([b,K.int_shape(x1)[3]])(x1)
    F = Lambda(forAll)([x1r])
    D = Dense(noClasse, activation='softmax',name='dense1')(F)
    extractor2 = Model(inputs=input, outputs=D)
    return extractor2

######### MAIN 

if "sess" in locals():     # this code prevent creating many intance of
    tf.compat.v1.reset_default_graph()
    sess.close()           # tensorflow in memory (that can cause crashes)
K.clear_session()
sess = tf.InteractiveSession()

model = Sequential()
model.add(Conv2D(filters=128, kernel_size=3,  strides=1, input_shape=(64,64,1),activation='relu'))
model.add(MaxPooling2D(pool_size=(2,2), strides=2))
model.add(Conv2D(filters=256, kernel_size=3, strides=1,   activation='relu'))
model.add(MaxPooling2D(pool_size=(2,2), strides=2))
model.add(Conv2D(filters=512, kernel_size=3,strides=1,activation='relu'))
model.add(MaxPooling2D(pool_size=(2,2), strides=2))

extract = compact(model,6,64,47)
extract.compile(optimizer=adam(lr=0.01,decay=0.00001), loss='categorical_crossentropy', metrics=['accuracy'])