Решение СЛАУ с использованием нейронной сети

model = tf.keras.models.Sequential()
 
model.add(tf.keras.layers.InputLayer(input_shape = X_train.shape[1:]))
model.add(tf.keras.layers.Dense(12, activation = 'sigmoid'))
model.add(tf.keras.layers.Dense(12, activation = 'sigmoid'))
model.add(tf.keras.layers.Dense(4, activation = 'relu'))
model.add(tf.keras.layers.Dense(n_outputs))
 
opt=tf.keras.optimizers.SGD(learning_rate = 0.001, momentum = 0.9)
 
model.compile(loss="mean_squared_error", optimizer=opt, metrics=['mae'])

class LinearSystemSolution(tf.keras.Model):
    def __init__(self, n_input_dimension, A, b, **kwargs):
        super(LinearSystemSolution, self).__init__(**kwargs)
        self.solution = tf.Variable(initial_value=tf.zeros((
                      n_input_dimension,1), dtype=tf.float64),
                      trainable= True, dtype=tf.float64, shape=(n_input_dimension,1))
        self.A = tf.convert_to_tensor(A)
        self.b = tf.convert_to_tensor(b)
    def call(self, x, training=True):
        error = tf.matmul(self.A, self.solution) - self.b
        self.add_loss(tf.reduce_sum(tf.square(error))/2)
        return self.solution
 
model = LinearSystemSolution(n_dim, A, b)
model.compile(optimizer = 'sgd')
 
history = model.fit(x_train, epochs = NET_EPOCHS, verbose=0, batch_size=NET_BATCH_SIZE)