Subclassing

import tensorflow as tf
from tensorflow.keras import layers, Model

# ===== Three ways to define a Keras model =====
# 1. Sequential   simplest; linear stack
# 2. Functional   most flexible; DAG with shared layers, multi-in/out
# 3. Subclass     full control; custom forward + training step

# ===== Subclass example =====
class MLP(Model):
    def __init__(self, hidden, num_classes):
        super().__init__()
        self.fc1 = layers.Dense(hidden, activation='relu')
        self.dropout = layers.Dropout(0.2)
        self.fc2 = layers.Dense(num_classes, activation='softmax')

    def call(self, x, training=False):
        h = self.fc1(x)
        h = self.dropout(h, training=training)
        return self.fc2(h)

model = MLP(hidden=64, num_classes=3)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Build by calling once (so it knows input shape):
import numpy as np
X = np.random.rand(100, 4); y = np.random.randint(0, 3, 100)
model.fit(X, y, epochs=2, verbose=0)

# ===== Custom training step =====
class CustomModel(Model):
    def __init__(self):
        super().__init__()
        self.dense = layers.Dense(3, activation='softmax')
    def call(self, x):
        return self.dense(x)
    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            preds = self(x, training=True)
            loss = self.compiled_loss(y, preds)
        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        self.compiled_metrics.update_state(y, preds)
        return {m.name: m.result() for m in self.metrics}

# Now model.fit() uses your train_step but you still benefit from .fit/.evaluate.

# ===== Custom layer =====
class GatedLinear(layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units
    def build(self, input_shape):
        self.w = self.add_weight('w', shape=(input_shape[-1], self.units),
                                 initializer='glorot_uniform')
        self.g = self.add_weight('g', shape=(input_shape[-1], self.units),
                                 initializer='glorot_uniform')
    def call(self, x):
        return tf.nn.sigmoid(x @ self.g) * (x @ self.w)

# ===== When to pick what =====
# Sequential   simple feed-forward
# Functional   anytime there are skip connections, multi-input/output, shared layers
# Subclass     when you need custom forward logic, custom train_step, or research code

# ===== Subclassing tradeoffs =====
# Pros: full control, easy debugging, dynamic shapes
# Cons: model.summary() less informative without explicit build_input_shape;
#       saving / loading requires custom config + get_config()

# ===== get_config + from_config (for serialisation) =====
class MLP2(Model):
    def __init__(self, hidden=64, num_classes=3, **kw):
        super().__init__(**kw)
        self.hidden = hidden; self.num_classes = num_classes
        self.fc1 = layers.Dense(hidden, activation='relu')
        self.fc2 = layers.Dense(num_classes, activation='softmax')
    def call(self, x): return self.fc2(self.fc1(x))
    def get_config(self):
        return { **super().get_config(),
                 'hidden': self.hidden, 'num_classes': self.num_classes }

# ===== Patterns to internalise =====
# - Default to Functional; reach for Subclass when you need custom logic
# - Implement build() if shapes depend on input
# - Override train_step for custom losses + manual gradient logic
# - get_config + register_keras_serializable for saving custom models

# ===== Pitfalls =====
# - Forgetting training=False at inference -> Dropout / BatchNorm misbehave
# - Subclassed models without get_config -> hard to deserialise
# - Defining layers inside call() -> new layers every call, never trained
# - Custom train_step that misses metric updates -> metrics never move

import tensorflow as tf
from tensorflow.keras import layers, Model
class MLP(Model):
    def __init__(self):
        super().__init__()
        self.d1 = layers.Dense(128, activation='relu')
        self.d2 = layers.Dense(10, activation='softmax')
    def call(self, x):
        return self.d2(self.d1(x))
model = MLP()