import tensorflow as tf
import tensorflow_datasets as tfds

(train_data, test_data), info = tfds.load(
    'mnist', split=['train', 'test'], as_supervised=True, with_info=True)

train_size = info.splits['train'].num_examples  # 60000
test_size = info.splits['test'].num_examples    # 10000

def normalize_img(image, label):
    return tf.cast(image, tf.float32) / 255.0, tf.cast(label, tf.int32)

train_data = train_data.map(normalize_img).shuffle(10000)
train_data = train_data.batch(128, drop_remainder=True)
test_data = test_data.map(normalize_img).batch(128, drop_remainder=True)
# 468 batches, each images shape (128, 28, 28, 1), labels shape (128,)

class CustomDense(tf.Module):
    def __init__(self, input_dim, output_dim, activation=tf.identity, name=None):
        super().__init__(name=name)
        self.W = tf.Variable(tf.random.normal([input_dim, output_dim], stddev=0.1))
        self.b = tf.Variable(tf.zeros([output_dim]))
        self.activation = activation

    def __call__(self, x):
        z = tf.matmul(x, self.W) + self.b
        return self.activation(z)

class MultiLayerNN(tf.Module):
    def __init__(self, input_dim, hidden_units=[128, 64], output_dim=10, activation=tf.nn.relu):
        super().__init__()
        self.flatten = tf.keras.layers.Flatten()
        self.hidden1 = CustomDense(input_dim, hidden_units[0], activation=activation)
        self.hidden2 = CustomDense(hidden_units[0], hidden_units[1], activation=activation)
        self.output_layer = CustomDense(hidden_units[1], output_dim, activation=tf.nn.softmax)

    def __call__(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = self.hidden2(x)
        return self.output_layer(x)

def compute_loss(y_pred, y_true):
    return tf.reduce_mean(
        tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)
    )

def train_model(model, dataset, learning_rate=0.01, epochs=5):
    optimizer = tf.optimizers.Adam(learning_rate)

    for epoch in range(epochs):
        total_loss, total_acc, num_batches = 0, 0, 0
        for x_batch, y_batch in dataset:
            with tf.GradientTape() as tape:
                y_pred = model(x_batch)
                loss = compute_loss(y_pred, y_batch)

            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

            predictions = tf.argmax(y_pred, axis=1, output_type=tf.int32)
            accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32)).numpy()
            total_loss += loss.numpy(); total_acc += accuracy; num_batches += 1

        print(f"Epoch {epoch+1}, Loss: {total_loss/num_batches:.4f}, Accuracy: {total_acc/num_batches:.4f}")

def evaluate_model(model, dataset):
    total_acc, num_batches = 0, 0
    for x_batch, y_batch in dataset:
        y_pred = model(x_batch)
        predictions = tf.argmax(y_pred, axis=1, output_type=tf.int32)
        total_acc += tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32)).numpy()
        num_batches += 1
    return total_acc / num_batches

mlp_model = MultiLayerNN(input_dim=784, hidden_units=[256, 128], output_dim=10)
train_model(mlp_model, train_data, learning_rate=0.001, epochs=5)
test_accuracy = evaluate_model(mlp_model, test_data)

Epoch 1, Loss: 1.5804, Accuracy: 0.8944
Epoch 5, Loss: 1.4851, Accuracy: 0.9778
Test Accuracy: 0.9732