Error description:

There is a phenomenon in which backpropagation is interrupted in a self-made layer called world_vec. When I checked the graph on tensorboard, the original layer was displayed as None and the connection with the subsequent layer was broken. What I expected was displayed in the input section of the original layer, but nothing was displayed in the output section.

UserWarning: Gradients do not exist for variables ['kernel', 'bias', 'kernel', 'bias', 'kernel', 'bias', 'gamma', 'beta', 'kernel', 'bias', 'gamma', 'beta', 'kernel', 'bias', 'gamma', 'beta', 'world_vec', 'world_vec'] when minimizing the loss. If using `model.compile()`, did you forget to provide a `loss` argument?

Also, when fitting the model, if I set the number of batches to 32, the following error will occur. It is fixed by setting the batch number to 1.

Arguments `target` and `output` must have the same shape up until the last dimension: target.shape=(32,), output.shape=(1, 2)
  File "C:Usersmnist.py", line 345, in <module>
    history = model.fit(x_train, y_train, batch_size=32, epochs=1, callbacks=[tensorboard_callback])
ValueError: Arguments `target` and `output` must have the same shape up until the last dimension: target.shape=(32,), output.shape=(1, 2)

–Explanation of the original layer–

Rather than building something that already exists, I came up with a completely new structure myself. What I created is what is called a world model. It simulates the state of the world using something like memory in LSTM, and then processes the memory with dense to generate an answer.

Original Layer Diagram

If a question gives us new knowledge or a hypothesis, we make changes to the memory (state of the world) accordingly and then simulate the hypothesis or knowledge. If there is a contradiction in the hypothesis, it can be detected. The following code experiments with the mnist dataset

original layer code

import tensorflow as tf
import numpy as np
import math


        
                   
class World_Vec(tf.keras.layers.Layer):
    def __init__(self, output_dim):
        super().__init__()
        self.output_dim = output_dim

        self.world_vec = None
        self.conditions_vec = None
        self.result_vec = None
        
        self.b1 = None
        self.b2 = None
        self.a = None

        self.conditions_vec_size = (10, 10, 10)
　　　　　　　　　　　　　　　　　　　　　
        self.world_vec_size = (10, 10, 10)
                                  


    def build(self, input_shape=(2, 10, 10, 10)):
        #The condition vector and result vector are in corresponding positions. It is difficult to explain in words, so if you did not understand, please see the diagram

        self.world_vec = self.add_weight(name="world_vec", shape=self.world_vec_size)
        self.conditions_vec = self.add_weight(name="conditions_vec", shape=self.conditions_vec_size)
        self.result_vec = self.add_weight(name="result_vec", shape=self.conditions_vec_size)

        self.b = self.add_weight(name="world_vec", shape=(self.conditions_vec_size[1], 1, 1))
        self.a = self.add_weight(name="world_vec", shape=(self.conditions_vec_size[1], 1, 1))

    

    
    def call(self, inputs, **kwargs):
        #Apply the hypotheses received in the input to the world vector
        #Example of input shape=(2, 10, 10, 10) input.shape[1]=number of hypotheses, input.shape[0]=two because we need a marker and a hypothesis to replace a particular world vector with a hypothesis. Example: if we want to hypothesize that if we drink water, we will get drunk, we need to change the result vector of getting drunk as a landmark to see where we can change the condition vector.
        inputs = tf.squeeze(inputs, axis=0)
        concatenated = tf.concat([self.world_vec, self.conditions_vec, self.result_vec], 0) #Consolidation to be completed in one go. I thought we could do them separately, but it would be inconvenient during softmax.

        #print("aiueo" + str(inputs.shape))

        concatenated, tekito1, tekito2 = self.change_world_vec(concatenated, inputs[0], inputs[1])
        #print("kakikukeko" + str(concatenated.shape))

        world_vec = concatenated[0:self.world_vec_size[0]]
        #print(self.world_vec.shape)
        conditions_vec = concatenated[self.world_vec_size[0]:self.world_vec_size[0] + self.world_vec_size[0]]
        #print(self.conditions_vec.shape)
        result_vec = concatenated[self.world_vec_size[0] + self.world_vec_size[0]:]
        

            
        
        for i in range(1):
            world_vec, conditions_vec, result_vec = self.change_world_vec(world_vec, conditions_vec, result_vec)

        return tf.reshape(tf.concat([world_vec, conditions_vec, result_vec], 0), shape=(1, -1))
        

        
    
    def softmax(self, logits):
        result = []
        sumed = np.sum(np.exp(logits))
        for i in logits.reshape((-1)):
            result.append(np.exp(i))

        result = np.array(result).reshape(logits.shape)

        return result
    
    
    def change_world_vec(self, world_vec, conditions_vec, result_vec):
        for i in range(self.world_vec.shape[0]):
            attention_weight = tf.constant([])
            for k in range(self.conditions_vec.shape[0]):
                
                world_vec_temp = tf.transpose(world_vec[i])
                conditions_vec_temp = conditions_vec[k]

                dot_temp = tf.tensordot(world_vec_temp, conditions_vec_temp, axes=1)
                dot_temp = tf.math.reduce_sum(dot_temp)
                #Measure the match between the condition vector and the world vector.
                
                normalized_dot = dot_temp/math.sqrt(world_vec.shape[-1]*world_vec.shape[-2])#normalize
                                       
                attention_weight = tf.keras.ops.append(attention_weight, normalized_dot)

            



            
            #Extracting from an result vector
            attention_weight = tf.nn.softmax(attention_weight)
            attention_weight = tf.reshape(attention_weight, (conditions_vec.shape[0], 1, 1))
            #print(conditions_vec.shape)
            taked_out_result_vec = result_vec * attention_weight
            taked_out_result_vec = tf.math.reduce_sum(taked_out_result_vec, axis=0)

            #Replace the world vector with the extracted one. First, delete the corresponding part of the world vector, then add it.
            erase_attention = tf.ones((conditions_vec.shape[0], 1, 1), dtype=tf.dtypes.float32)
            erase_attention = erase_attention - tf.math.sigmoid(self.a * (attention_weight-self.b) + self.b)#Think of self.a(attention_weight-self.b)+self.b as normalization.
            erase_attention = tf.math.reduce_sum(erase_attention)

           

            erased_world_vec = world_vec[i] * attention_weight
            if i == 0:
                changed_world_vec = world_vec[i] + taked_out_result_vec

            else:
                changed_world_vec = tf.concat([changed_world_vec, world_vec[i] + taked_out_result_vec], axis=0)
        
        return world_vec, conditions_vec, result_vec
    
    

keras functional api code

    h = keras.layers.Dense(2000, activation="relu")(h)
    h = keras.layers.BatchNormalization(
        axis=-1,
        momentum=0.99,
        epsilon=0.001,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_initializer="ones",
        moving_mean_initializer="zeros",
        moving_variance_initializer="ones",
        beta_regularizer=None,
        gamma_regularizer=None,
        beta_constraint=None,
        gamma_constraint=None,
        synchronized=False)(h)
    
    h = keras.layers.Reshape((2, 10, 10, 10))(h)
    
    all_vec = World_Vec(4)(h)

    h1 = keras.layers.Dense(64, activation="relu")(all_vec)
    h1 = keras.layers.BatchNormalization(
        axis=-1,
        momentum=0.99,
        epsilon=0.001,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_initializer="ones",
        moving_mean_initializer="zeros",
        moving_variance_initializer="ones",
        beta_regularizer=None,
        gamma_regularizer=None,
        beta_constraint=None,
        gamma_constraint=None,
        synchronized=False)(h1)


    #h = keras.layers.Dropout(dropout_rate)(h)
    h1 = keras.layers.Dense(64, activation="relu")(h1)

    h1 = keras.layers.BatchNormalization(
        axis=-1,
        momentum=0.99,
        epsilon=0.001,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_initializer="ones",
        moving_mean_initializer="zeros",
        moving_variance_initializer="ones",
        beta_regularizer=None,
        gamma_regularizer=None,
        beta_constraint=None,
        gamma_constraint=None,
        synchronized=False)(h1)
    
    h1 = keras.layers.Dense(64, activation="relu")(h1)
    h1 = keras.layers.BatchNormalization(
        axis=-1,
        momentum=0.99,
        epsilon=0.001,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_initializer="ones",
        moving_mean_initializer="zeros",
        moving_variance_initializer="ones",
        beta_regularizer=None,
        gamma_regularizer=None,
        beta_constraint=None,
        gamma_constraint=None,
        synchronized=False)(h1)
    outputs = keras.layers.Dense(2, activation="relu")(h1)

    model = keras.Model(inputs=inputs, outputs=outputs, name="mnist_model")
    model.summary()


if __name__ == "__main__":
    (x_train, y_train_temp), (x_test, y_test_temp) = keras.datasets.mnist.load_data()
    print(x_train.shape)
    x_train = x_train.reshape(-1, 28, 28, 1).astype("float32") / 255
    x_test = x_test.reshape(-1, 28, 28, 1).astype("float32") / 255

    print(y_train_temp[1:20])

    y_train = []

    for i in y_train_temp:
        if i == 5:
            y_train.append(1)
            
            
        else:
            y_train.append(0)
            
    y_train = np.array(y_train).reshape(-1, )
    y_test = []

    for i in y_test_temp:
        if i == 5:
            y_test.append(1)
            
        else:
            y_test.append(0)
    y_test = np.array(y_test).reshape(-1, )


    print(y_test.shape)

    model = make_model()
    model.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer=keras.optimizers.Adam(),
        metrics=["accuracy"]
    )

    logdir = "/logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")

    tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

    history = model.fit(x_train, y_train, batch_size=1, epochs=1, callbacks=[tensorboard_callback])

The state of the world is expressed in terms of weights, but I thought it was wrong to mess with the weights, so I copied the weights to another variable and then did it.