input: 
  import tensorflow as tf
  
  # Import MNIST data
  from tensorflow.examples.tutorials.mnist import input_data
  mnist = input_data.read_data_sets("/tmp/data", one_hot=True)
  
  def train_mnist_model(learning_rate: float, training_epochs: int, batch_size: int) -> None:
      """ Trains a neural network model for classifying hand-written digits from the MNIST dataset.
      The model consists of two hidden layers with 256 neurons each and an output layer with 10 neurons 
      corresponding to the 10 digit classes. The model is trained using the Adam optimizer and 
      softmax cross-entropy loss.
      
      Example:
      for train_mnist_model(0.001, 15, 100)
      Epoch: 0001 cost= 1.183523522
      Epoch: 0002 cost= 0.441088046
      Epoch: 0003 cost= 0.337134511
      ...
      Optimization Finished!
      Accuracy: 0.92
      """
      
      # Network Parameters
      n_hidden_1 = 256 # 1st layer number of neurons
      n_hidden_2 = 256 # 2nd layer number of neurons
      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("float", [None, n_input])
      Y = tf.placeholder("float", [None, n_classes])
  
      # Store layers weight & bias
      weights = {
          'h1': tf.Variable(tf.random_normal([n_input, n_hidden_1])),
          'h2': tf.Variable(tf.random_normal([n_hidden_1, n_hidden_2])),
          'out': tf.Variable(tf.random_normal([n_hidden_2, n_classes]))
      }
      biases = {
          'b1': tf.Variable(tf.random_normal([n_hidden_1])),
          'b2': tf.Variable(tf.random_normal([n_hidden_2])),
          'out': tf.Variable(tf.random_normal([n_classes]))
      }
  
      # Create model
      def multilayer_perceptron(x):
          layer_1 = tf.add(tf.matmul(x, weights['h1']), biases['b1'])
          layer_2 = tf.add(tf.matmul(layer_1, weights['h2']), biases['b2'])
          out_layer = tf.matmul(layer_2, weights['out']) + biases['out']
          return out_layer
  
      # Construct model
      logits = multilayer_perceptron(X)
  
      # Define loss and optimizer
      loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
          logits=logits, labels=Y))
      optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
      train_op = optimizer.minimize(loss_op)
  
      # Initializing the variables
      init = tf.global_variables_initializer()
  
      # Train model
      with tf.Session() as model:
          model.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)
                  _, c = model.run([train_op, loss_op], feed_dict={X: batch_x, Y: batch_y})
                  avg_cost += c / total_batch
              if epoch % 1 == 0:
                  print("Epoch:", '%04d' % (epoch+1), "cost=", "{:.9f}".format(avg_cost))
          print("Optimization Finished!")
  
          # Test model
          pred = tf.nn.softmax(logits)  # Apply softmax to logits
          correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(Y, 1))
          accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
          print("Accuracy:", accuracy.eval({X: mnist.test.images, Y: mnist.test.labels}))
  
output:     
          # Train model
          with tf.Session() as model:
              model.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)
                      _, c = model.run([train_op, loss_op], feed_dict={X: batch_x, Y: batch_y})
                      avg_cost += c / total_batch
                  if epoch % 1 == 0:
                      print("Epoch:", '%04d' % (epoch+1), "cost=", "{:.9f}".format(avg_cost))
              print("Optimization Finished!")
  
              # Test model
              pred = tf.nn.softmax(logits)  # Apply softmax to logits
              correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(Y, 1))
              accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
              print("Accuracy:", accuracy.eval({X: mnist.test.images, Y: mnist.test.labels}))
  