input: 
  def train_mnist_nn_model(mnist: 'MNISTData', learning_rate: float, training_epochs: int, batch_size: int, display_step: int) -> float:
      """ Train a neural network model to classify hand-written digits from the MNIST dataset and return the accuracy on the test set.
      >>> mnist = input_data.read_data_sets("/tmp/data", one_hot=True)
      >>> accuracy = train_mnist_nn_model(mnist, 0.001, 15, 100, 1)
      >>> print("Final Test Accuracy:", accuracy)
      Final Test Accuracy: 0.95
      """
  
output:     
      # Parameters
      n_input = 784  # MNIST data input (img shape: 28*28)
      n_classes = 10  # MNIST total classes (0-9 digits)
  
      # tf Graph input
      x = tf.placeholder(tf.float32, [None, n_input])
      y = tf.placeholder(tf.float32, [None, n_classes])
  
      # Create model
      def multilayer_perceptron(x, weights, biases):
          layer_1 = tf.add(tf.matmul(x, weights['h1']), biases['b1'])
          layer_1 = tf.nn.relu(layer_1)
          out_layer = tf.matmul(layer_1, weights['out']) + biases['out']
          return out_layer
  
      # Store layers weight & bias
      weights = {
          'h1': tf.Variable(tf.random_normal([n_input, 256])),
          'out': tf.Variable(tf.random_normal([256, n_classes]))
      }
      biases = {
          'b1': tf.Variable(tf.random_normal([256])),
          'out': tf.Variable(tf.random_normal([n_classes]))
      }
  
      # Construct model
      pred = multilayer_perceptron(x, weights, biases)
  
      # Define loss and optimizer
      cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
      optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)
  
      # Initializing the variables
      init = tf.global_variables_initializer()
  
      # Launch the graph
      with tf.Session() as sess:
          sess.run(init)
  
          # Training cycle
          for epoch in range(training_epochs):
              avg_cost = 0.
              total_batch = int(mnist.train.num_examples/batch_size)
              # Loop over all batches
              for i in range(total_batch):
                  batch_x, batch_y = mnist.train.next_batch(batch_size)
                  # Run optimization op (backprop) and cost op (to get loss value)
                  _, c = sess.run([optimizer, cost], feed_dict={x: batch_x, y: batch_y})
                  # Compute average loss
                  avg_cost += c / total_batch
              # Display logs per epoch step
              if epoch % display_step == 0:
                  print(f"Epoch: {epoch+1} cost={avg_cost}")
  
          print("Optimization Finished!")
  
          # Test model
          correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
          # Calculate accuracy
          accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
          test_accuracy = accuracy.eval({x: mnist.test.images, y: mnist.test.labels})
          return test_accuracy
  