deeplearn_numpy.py

# deeplearn_numpy.py
# NumPy-only, parametric deep learning mini-framework in a single file.
# - Supports configurable layers via registries and a JSON-like config
# - Core layers: Dense, Conv2D (im2col), MaxPool2D, AvgPool2D, Flatten, Dropout, Activations, BatchNorm (1D/2D)
# - Losses: MSE, CrossEntropy (with stable softmax), BinaryCrossEntropy
# - Optimizers: SGD, Momentum, Adam
# - Model: Sequential builder from config; training loop (fit/evaluate/predict)
# - Regularization: L1/L2
# - Serialization: save/load weights; save/load config
# - Metrics: accuracy (multi-class), binary_accuracy, top_k_accuracy
#
# Usage (example at bottom of file):
#   from deeplearn_numpy import Model, build_model_from_config
#   model = build_model_from_config(config)
#   model.fit(x_train, y_train, epochs=5, batch_size=32)
#
# This file is self-contained and uses only NumPy.

from __future__ import annotations
import numpy as np
from typing import Any, Dict, List, Tuple, Optional, Iterable, Callable

# ------------------------------------------------------------
# Utility: RNG and dtypes
# ------------------------------------------------------------

def _get_rng(seed: Optional[int]) -> np.random.Generator:
    return np.random.default_rng(seed) if seed is not None else np.random.default_rng()

DEFAULT_DTYPE = np.float32

# ------------------------------------------------------------
# Parameter wrapper
# ------------------------------------------------------------

class Parameter:
    def __init__(self, data: np.ndarray, name: Optional[str] = None):
        self.data = data.astype(DEFAULT_DTYPE, copy=False)
        self.grad = np.zeros_like(self.data)
        self.name = name

    def zero_grad(self):
        self.grad[...] = 0

# ------------------------------------------------------------
# Initializers (registry + functions)
# ------------------------------------------------------------

INITIALIZER_REGISTRY: Dict[str, Callable[[Tuple[int, ...], Optional[int]], np.ndarray]] = {}

def register_initializer(name: str):
    def deco(fn):
        INITIALIZER_REGISTRY[name] = fn
        return fn
    return deco

@register_initializer("zeros")
def init_zeros(shape: Tuple[int, ...], seed: Optional[int] = None):
    return np.zeros(shape, dtype=DEFAULT_DTYPE)

@register_initializer("ones")
def init_ones(shape: Tuple[int, ...], seed: Optional[int] = None):
    return np.ones(shape, dtype=DEFAULT_DTYPE)

@register_initializer("random_normal")
def init_random_normal(shape: Tuple[int, ...], seed: Optional[int] = None):
    rng = _get_rng(seed)
    return rng.standard_normal(size=shape).astype(DEFAULT_DTYPE)

@register_initializer("glorot_uniform")
def init_glorot_uniform(shape: Tuple[int, ...], seed: Optional[int] = None):
    # fan_in, fan_out heuristic (for Dense/Conv2D weights)
    if len(shape) == 2:
        fan_in, fan_out = shape[0], shape[1]
    elif len(shape) == 4:  # (out_ch, in_ch, kH, kW)
        fan_in = shape[1] * shape[2] * shape[3]
        fan_out = shape[0] * shape[2] * shape[3]
    else:
        fan_in = int(np.prod(shape[:-1])) if len(shape) > 1 else shape[0]
        fan_out = shape[-1]
    limit = np.sqrt(6.0 / (fan_in + fan_out))
    rng = _get_rng(seed)
    return rng.uniform(-limit, limit, size=shape).astype(DEFAULT_DTYPE)

@register_initializer("he_normal")
def init_he_normal(shape: Tuple[int, ...], seed: Optional[int] = None):
    # Good for ReLU
    if len(shape) == 2:
        fan_in = shape[0]
    elif len(shape) == 4:
        fan_in = shape[1] * shape[2] * shape[3]
    else:
        fan_in = int(np.prod(shape[:-1])) if len(shape) > 1 else shape[0]
    std = np.sqrt(2.0 / fan_in)
    rng = _get_rng(seed)
    return (rng.standard_normal(size=shape) * std).astype(DEFAULT_DTYPE)

def get_initializer(name: str) -> Callable[[Tuple[int, ...], Optional[int]], np.ndarray]:
    if name not in INITIALIZER_REGISTRY:
        raise ValueError(f"Unknown initializer: {name}")
    return INITIALIZER_REGISTRY[name]

# ------------------------------------------------------------
# Activations as layers (registry-based)
# ------------------------------------------------------------

class Layer:
    def __init__(self):
        self.built = False
        self.training = True  # toggled by model
        self._params: List[Parameter] = []
        self._cache: Dict[str, Any] = {}
        self.name: Optional[str] = None

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        self.built = True
        return input_shape

    def forward(self, x: np.ndarray) -> np.ndarray:
        raise NotImplementedError

    def backward(self, dy: np.ndarray) -> np.ndarray:
        raise NotImplementedError

    def params(self) -> Iterable[Parameter]:
        return self._params

    def zero_grads(self):
        for p in self._params:
            p.zero_grad()

    def set_training(self, mode: bool):
        self.training = mode

    # Optional: for shape inference without allocating params
    def infer_output_shape(self, input_shape: Tuple[int, ...]) -> Tuple[int, ...]:
        # default: identity
        return input_shape

# Activation registry
ACTIVATION_REGISTRY: Dict[str, "Layer"] = {}
LAYER_REGISTRY: Dict[str, Callable[..., Layer]] = {}

def register_layer(name: str):
    def deco(cls):
        LAYER_REGISTRY[name] = cls
        return cls
    return deco

def register_activation(name: str):
    def deco(cls):
        ACTIVATION_REGISTRY[name] = cls
        return cls
    return deco

# ---- Activation implementations ----

@register_activation("relu")
@register_layer("ReLU")
class ReLU(Layer):
    def forward(self, x):
        self._cache["mask"] = x > 0
        return np.where(self._cache["mask"], x, 0).astype(DEFAULT_DTYPE)

    def backward(self, dy):
        return dy * self._cache["mask"]

@register_activation("sigmoid")
@register_layer("Sigmoid")
class Sigmoid(Layer):
    def forward(self, x):
        out = 1.0 / (1.0 + np.exp(-x))
        self._cache["out"] = out
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        out = self._cache["out"]
        return dy * out * (1.0 - out)

@register_activation("tanh")
@register_layer("Tanh")
class Tanh(Layer):
    def forward(self, x):
        out = np.tanh(x)
        self._cache["out"] = out
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        out = self._cache["out"]
        return dy * (1.0 - out ** 2)

@register_activation("leaky_relu")
@register_layer("LeakyReLU")
class LeakyReLU(Layer):
    def __init__(self, alpha: float = 0.01):
        super().__init__()
        self.alpha = alpha

    def forward(self, x):
        self._cache["x"] = x
        return np.where(x > 0, x, self.alpha * x).astype(DEFAULT_DTYPE)

    def backward(self, dy):
        x = self._cache["x"]
        dx = np.ones_like(x)
        dx[x < 0] = self.alpha
        return dy * dx

@register_activation("softmax")
@register_layer("Softmax")
class Softmax(Layer):
    def __init__(self, axis: int = -1):
        super().__init__()
        self.axis = axis

    def forward(self, x):
        # stable softmax
        x_shift = x - np.max(x, axis=self.axis, keepdims=True)
        exps = np.exp(x_shift)
        sums = np.sum(exps, axis=self.axis, keepdims=True)
        out = exps / sums
        self._cache["out"] = out
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        # General softmax backward: J = diag(s) - s s^T; dy @ J
        s = self._cache["out"]
        # Works for 2D (batch, classes) or higher dims along axis
        # Flatten axis to -1 for computation
        orig_shape = s.shape
        axis = self.axis if self.axis >= 0 else s.ndim + self.axis
        s2d = s.reshape(-1, orig_shape[axis])
        dy2d = dy.reshape(-1, orig_shape[axis])
        # gradient per row: dy - sum(dy*s) * s
        dot = np.sum(dy2d * s2d, axis=1, keepdims=True)
        dx2d = s2d * (dy2d - dot)
        return dx2d.reshape(orig_shape)

# ------------------------------------------------------------
# Core layers
# ------------------------------------------------------------

@register_layer("Dense")
class Dense(Layer):
    def __init__(self, units: int, use_bias: bool = True,
                 kernel_initializer: str = "glorot_uniform",
                 bias_initializer: str = "zeros",
                 seed: Optional[int] = None):
        super().__init__()
        self.units = units
        self.use_bias = use_bias
        self.kernel_initializer = kernel_initializer
        self.bias_initializer = bias_initializer
        self.seed = seed

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        assert len(input_shape) == 2, "Dense expects (batch, features)"
        in_dim = input_shape[1]
        W = get_initializer(self.kernel_initializer)((in_dim, self.units), seed=self.seed)
        self.W = Parameter(W, name="kernel")
        self._params.append(self.W)
        if self.use_bias:
            b = get_initializer(self.bias_initializer)((self.units,), seed=self.seed)
            self.b = Parameter(b, name="bias")
            self._params.append(self.b)
        else:
            self.b = None
        self.built = True
        return (input_shape[0], self.units)

    def forward(self, x):
        self._cache["x"] = x
        out = x @ self.W.data
        if self.b is not None:
            out = out + self.b.data
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        x = self._cache["x"]
        # grads
        self.W.grad += x.T @ dy
        if self.b is not None:
            self.b.grad += np.sum(dy, axis=0)
        dx = dy @ self.W.data.T
        return dx

# ---- im2col/col2im utilities ----

def get_output_dim(in_size, k, stride, pad, dilation):
    # For 'valid' pad: pad=0; for 'same': computed elsewhere
    return (in_size + 2*pad - dilation*(k-1) - 1)//stride + 1

def compute_same_padding(in_size, k, stride, dilation):
    # Output size ceil(in/stride); padding to make that
    out_size = int(np.ceil(in_size / stride))
    pad_needed = max(0, (out_size - 1)*stride + dilation*(k - 1) + 1 - in_size)
    # split pad into left/right (we'll use symmetric here)
    pad_left = pad_needed // 2
    pad_right = pad_needed - pad_left
    return pad_left, pad_right, out_size

def im2col_indices(x, kH, kW, pad_h, pad_w, stride_h, stride_w, dilation_h=1, dilation_w=1):
    # x: (N, C, H, W)
    N, C, H, W = x.shape
    eff_kH = dilation_h*(kH-1)+1
    eff_kW = dilation_w*(kW-1)+1
    out_h = (H + pad_h*2 - eff_kH)//stride_h + 1
    out_w = (W + pad_w*2 - eff_kW)//stride_w + 1

    x_padded = np.pad(x, ((0,0),(0,0),(pad_h,pad_h),(pad_w,pad_w)), mode='constant')

    i0 = np.repeat(np.arange(kH), kW)
    i0 = np.tile(i0, C)
    i1 = stride_h * np.repeat(np.arange(out_h), out_w)
    j0 = np.tile(np.arange(kW), kH * C)
    j1 = stride_w * np.tile(np.arange(out_w), out_h)
    i = i0.reshape(-1,1) * dilation_h + i1.reshape(1,-1)
    j = j0.reshape(-1,1) * dilation_w + j1.reshape(1,-1)

    k = np.repeat(np.arange(C), kH * kW).reshape(-1,1)

    cols = x_padded[:, k, i, j]  # (N, C*kH*kW, out_h*out_w)
    cols = cols.transpose(1,2,0).reshape(C*kH*kW, N*out_h*out_w)
    return cols, out_h, out_w

def col2im_indices(cols, x_shape, kH, kW, pad_h, pad_w, stride_h, stride_w, dilation_h=1, dilation_w=1):
    N, C, H, W = x_shape
    eff_kH = dilation_h*(kH-1)+1
    eff_kW = dilation_w*(kW-1)+1
    out_h = (H + pad_h*2 - eff_kH)//stride_h + 1
    out_w = (W + pad_w*2 - eff_kW)//stride_w + 1

    cols_reshaped = cols.reshape(C*kH*kW, N, out_h*out_w).transpose(1,0,2)
    x_padded = np.zeros((N, C, H + 2*pad_h, W + 2*pad_w), dtype=cols.dtype)

    i0 = np.repeat(np.arange(kH), kW)
    i0 = np.tile(i0, C)
    i1 = stride_h * np.repeat(np.arange(out_h), out_w)
    j0 = np.tile(np.arange(kW), kH * C)
    j1 = stride_w * np.tile(np.arange(out_w), out_h)
    i = i0.reshape(-1,1) * dilation_h + i1.reshape(1,-1)
    j = j0.reshape(-1,1) * dilation_w + j1.reshape(1,-1)
    k = np.repeat(np.arange(C), kH * kW).reshape(-1,1)

    for n in range(N):
        x_padded[n, k, i, j] += cols_reshaped[n]

    return x_padded[:, :, pad_h:pad_h+H, pad_w:pad_w+W]

@register_layer("Conv2D")
class Conv2D(Layer):
    def __init__(self, filters: int, kernel_size: int | Tuple[int,int],
                 stride: int | Tuple[int,int] = 1,
                 padding: str | int | Tuple[int,int] = "valid",
                 dilation: int | Tuple[int,int] = 1,
                 use_bias: bool = True,
                 kernel_initializer: str = "he_normal",
                 bias_initializer: str = "zeros",
                 seed: Optional[int] = None):
        super().__init__()
        self.filters = filters
        if isinstance(kernel_size, int):
            self.kH = self.kW = kernel_size
        else:
            self.kH, self.kW = kernel_size
        if isinstance(stride, int):
            self.sH = self.sW = stride
        else:
            self.sH, self.sW = stride
        self.padding = padding
        if isinstance(dilation, int):
            self.dH = self.dW = dilation
        else:
            self.dH, self.dW = dilation
        self.use_bias = use_bias
        self.kernel_initializer = kernel_initializer
        self.bias_initializer = bias_initializer
        self.seed = seed

    def _compute_padding(self, H, W):
        if self.padding == "valid":
            return (0,0,0,0)
        if self.padding == "same":
            ph1, ph2, out_h = compute_same_padding(H, self.kH, self.sH, self.dH)
            pw1, pw2, out_w = compute_same_padding(W, self.kW, self.sW, self.dW)
            # symmetric padding; using total pads
            return (ph1, pw1, ph2, pw2)
        if isinstance(self.padding, int):
            return (self.padding, self.padding, self.padding, self.padding)
        if isinstance(self.padding, tuple):
            pH, pW = self.padding
            return (pH, pW, pH, pW)
        raise ValueError("Unsupported padding")

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        # input: (N, C, H, W)
        assert len(input_shape) == 4, "Conv2D expects (N, C, H, W)"
        N, C, H, W = input_shape
        self.in_channels = C

        pad_top, pad_left, pad_bottom, pad_right = self._compute_padding(H, W)
        eff_kH = self.dH*(self.kH-1)+1
        eff_kW = self.dW*(self.kW-1)+1
        out_h = (H + pad_top + pad_bottom - eff_kH)//self.sH + 1
        out_w = (W + pad_left + pad_right - eff_kW)//self.sW + 1
        Wshape = (self.filters, C, self.kH, self.kW)
        W = get_initializer(self.kernel_initializer)(Wshape, seed=self.seed)
        self.W = Parameter(W, name="kernel")
        self._params.append(self.W)
        if self.use_bias:
            b = get_initializer(self.bias_initializer)((self.filters,), seed=self.seed)
            self.b = Parameter(b, name="bias")
            self._params.append(self.b)
        else:
            self.b = None
        self.pad = (pad_top, pad_left, pad_bottom, pad_right)
        self.built = True
        return (N, self.filters, out_h, out_w)

    def forward(self, x):
        self._cache["x_shape"] = x.shape
        pt, pl, pb, pr = self.pad
        cols, out_h, out_w = im2col_indices(
            x, self.kH, self.kW, pt, pl, self.sH, self.sW, self.dH, self.dW
        )
        W_col = self.W.data.reshape(self.filters, -1)
        out = W_col @ cols  # (filters, N*out_h*out_w)
        if self.b is not None:
            out += self.b.data[:, None]
        N = x.shape[0]
        out = out.reshape(self.filters, -1, N).transpose(2, 0, 1)
        out = out.reshape(N, self.filters, out_h, out_w)
        self._cache["cols"] = cols
        self._cache["out_shape"] = (N, self.filters, out_h, out_w)
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        cols = self._cache["cols"]
        N, F, out_h, out_w = self._cache["out_shape"]
        W_col = self.W.data.reshape(self.filters, -1)

        dy_reshaped = dy.reshape(N, F, out_h*out_w).transpose(1,2,0).reshape(F, -1)
        # gradients
        self.W.grad += (dy_reshaped @ cols.T).reshape(self.W.data.shape)
        if self.b is not None:
            self.b.grad += np.sum(dy_reshaped, axis=1)

        dcols = W_col.T @ dy_reshaped
        pt, pl, pb, pr = self.pad
        dx = col2im_indices(
            dcols, self._cache["x_shape"], self.kH, self.kW, pt, pl, self.sH, self.sW, self.dH, self.dW
        )
        return dx

@register_layer("MaxPool2D")
class MaxPool2D(Layer):
    def __init__(self, pool_size: int | Tuple[int,int] = 2, stride: Optional[int | Tuple[int,int]] = None, padding: str | int | Tuple[int,int] = "valid"):
        super().__init__()
        if isinstance(pool_size, int):
            self.kH = self.kW = pool_size
        else:
            self.kH, self.kW = pool_size
        if stride is None:
            self.sH = self.kH
            self.sW = self.kW
        elif isinstance(stride, int):
            self.sH = self.sW = stride
        else:
            self.sH, self.sW = stride
        self.padding = padding

    def _compute_padding(self, H, W):
        if self.padding == "valid":
            return (0,0,0,0)
        if self.padding == "same":
            ph1, ph2, _ = compute_same_padding(H, self.kH, self.sH, 1)
            pw1, pw2, _ = compute_same_padding(W, self.kW, self.sW, 1)
            return (ph1, pw1, ph2, pw2)
        if isinstance(self.padding, int):
            return (self.padding, self.padding, self.padding, self.padding)
        if isinstance(self.padding, tuple):
            pH, pW = self.padding
            return (pH, pW, pH, pW)
        raise ValueError("Unsupported padding")

    def build(self, input_shape, **kwargs):
        N, C, H, W = input_shape
        pt, pl, pb, pr = self._compute_padding(H, W)
        out_h = (H + pt + pb - self.kH)//self.sH + 1
        out_w = (W + pl + pr - self.kW)//self.sW + 1
        self.pad = (pt, pl, pb, pr)
        self.built = True
        return (N, C, out_h, out_w)

    def forward(self, x):
        N, C, H, W = x.shape
        pt, pl, pb, pr = self.pad
        x_padded = np.pad(x, ((0,0),(0,0),(pt,pb),(pl,pr)), mode='constant', constant_values=-np.inf)
        out_h = (H + pt + pb - self.kH)//self.sH + 1
        out_w = (W + pl + pr - self.kW)//self.sW + 1

        self._cache["x_shape"] = x.shape
        self._cache["pad"] = self.pad

        out = np.empty((N, C, out_h, out_w), dtype=DEFAULT_DTYPE)
        self._cache["max_idx"] = np.zeros_like(out, dtype=np.int32)

        for i in range(out_h):
            for j in range(out_w):
                h0 = i*self.sH
                w0 = j*self.sW
                window = x_padded[:, :, h0:h0+self.kH, w0:w0+self.kW]  # (N, C, kH, kW)
                flat = window.reshape(N, C, -1)
                idx = np.argmax(flat, axis=2)
                self._cache["max_idx"][:, :, i, j] = idx
                out[:, :, i, j] = np.take_along_axis(flat, idx[..., None], axis=2).squeeze(-1)
        return out

    def backward(self, dy):
        N, C, H, W = self._cache["x_shape"]
        pt, pl, pb, pr = self._cache["pad"]
        out_h, out_w = dy.shape[2], dy.shape[3]
        dx_padded = np.zeros((N, C, H + pt + pb, W + pl + pr), dtype=DEFAULT_DTYPE)
        idx = self._cache["max_idx"]
        for i in range(out_h):
            for j in range(out_w):
                h0 = i*self.sH
                w0 = j*self.sW
                mask = np.zeros((N, C, self.kH*self.kW), dtype=DEFAULT_DTYPE)
                flat_idx = idx[:, :, i, j]
                mask.reshape(N*C, -1)[np.arange(N*C), flat_idx.reshape(-1)] = dy[:, :, i, j].reshape(-1)
                dx_padded[:, :, h0:h0+self.kH, w0:w0+self.kW] += mask.reshape(N, C, self.kH, self.kW)
        return dx_padded[:, :, pt:pt+H, pl:pl+W]

@register_layer("AvgPool2D")
class AvgPool2D(Layer):
    def __init__(self, pool_size: int | Tuple[int,int] = 2, stride: Optional[int | Tuple[int,int]] = None, padding: str | int | Tuple[int,int] = "valid"):
        super().__init__()
        if isinstance(pool_size, int):
            self.kH = self.kW = pool_size
        else:
            self.kH, self.kW = pool_size
        if stride is None:
            self.sH = self.kH
            self.sW = self.kW
        elif isinstance(stride, int):
            self.sH = self.sW = stride
        else:
            self.sH, self.sW = stride
        self.padding = padding

    def _compute_padding(self, H, W):
        if self.padding == "valid":
            return (0,0,0,0)
        if self.padding == "same":
            ph1, ph2, _ = compute_same_padding(H, self.kH, self.sH, 1)
            pw1, pw2, _ = compute_same_padding(W, self.kW, self.sW, 1)
            return (ph1, pw1, ph2, pw2)
        if isinstance(self.padding, int):
            return (self.padding, self.padding, self.padding, self.padding)
        if isinstance(self.padding, tuple):
            pH, pW = self.padding
            return (pH, pW, pH, pW)
        raise ValueError("Unsupported padding")

    def build(self, input_shape, **kwargs):
        N, C, H, W = input_shape
        pt, pl, pb, pr = self._compute_padding(H, W)
        out_h = (H + pt + pb - self.kH)//self.sH + 1
        out_w = (W + pl + pr - self.kW)//self.sW + 1
        self.pad = (pt, pl, pb, pr)
        self.built = True
        return (N, C, out_h, out_w)

    def forward(self, x):
        N, C, H, W = x.shape
        pt, pl, pb, pr = self.pad
        x_padded = np.pad(x, ((0,0),(0,0),(pt,pb),(pl,pr)), mode='constant', constant_values=0.0)
        out_h = (H + pt + pb - self.kH)//self.sH + 1
        out_w = (W + pl + pr - self.kW)//self.sW + 1
        out = np.empty((N, C, out_h, out_w), dtype=DEFAULT_DTYPE)

        for i in range(out_h):
            for j in range(out_w):
                h0 = i*self.sH
                w0 = j*self.sW
                window = x_padded[:, :, h0:h0+self.kH, w0:w0+self.kW]
                out[:, :, i, j] = np.mean(window, axis=(2,3))
        self._cache["x_shape"] = x.shape
        self._cache["pad"] = self.pad
        return out

    def backward(self, dy):
        N, C, H, W = self._cache["x_shape"]
        pt, pl, pb, pr = self._cache["pad"]
        out_h, out_w = dy.shape[2], dy.shape[3]
        dx_padded = np.zeros((N, C, H + pt + pb, W + pl + pr), dtype=DEFAULT_DTYPE)
        scale = 1.0 / (self.kH * self.kW)
        for i in range(out_h):
            for j in range(out_w):
                h0 = i*self.sH
                w0 = j*self.sW
                dx_padded[:, :, h0:h0+self.kH, w0:w0+self.kW] += dy[:, :, i, j][:, :, None, None] * scale
        return dx_padded[:, :, pt:pt+H, pl:pl+W]

@register_layer("Reshape")
class Reshape(Layer):
    def __init__(self, target_shape: Tuple[int, ...]):
        super().__init__()
        self.target_shape = tuple(target_shape)
        self.input_shape = None

    def build(self, input_shape, **kwargs):
        self.input_shape = input_shape[1:]
        return (input_shape[0],) + self.target_shape

    def forward(self, x):
        self._cache["in_shape"] = x.shape
        return x.reshape((x.shape[0],) + self.target_shape)

    def backward(self, dy):
        return dy.reshape(self._cache["in_shape"])

@register_layer("UpSample2D")
class UpSample2D(Layer):
    def __init__(self, scale: int = 2):
        super().__init__()
        self.scale = scale

    def build(self, input_shape, **kwargs):
        N, C, H, W = input_shape
        return (N, C, H * self.scale, W * self.scale)

    def forward(self, x):
        self._cache["in_shape"] = x.shape
        return np.repeat(np.repeat(x, self.scale, axis=2), self.scale, axis=3)

    def backward(self, dy):
        scale = self.scale
        return dy.reshape(
            dy.shape[0], dy.shape[1],
            dy.shape[2]//scale, scale,
            dy.shape[3]//scale, scale
        ).mean(axis=(3,5))

@register_layer("LayerNorm")
class LayerNorm(Layer):
    def __init__(self, eps: float = 1e-5, affine: bool = True):
        super().__init__()
        self.eps = eps
        self.affine = affine

    def build(self, input_shape, **kwargs):
        self.normalized_shape = input_shape[1:]
        if self.affine:
            self.gamma = Parameter(np.ones(self.normalized_shape, dtype=DEFAULT_DTYPE), "gamma")
            self.beta = Parameter(np.zeros(self.normalized_shape, dtype=DEFAULT_DTYPE), "beta")
            self._params += [self.gamma, self.beta]
        return input_shape

    def forward(self, x):
        mean = np.mean(x, axis=1, keepdims=True)
        var = np.var(x, axis=1, keepdims=True)
        self._cache["x_centered"] = x - mean
        self._cache["inv_std"] = 1.0 / np.sqrt(var + self.eps)
        x_hat = self._cache["x_centered"] * self._cache["inv_std"]
        if self.affine:
            return x_hat * self.gamma.data + self.beta.data
        return x_hat

    def backward(self, dy):
        x_centered = self._cache["x_centered"]
        inv_std = self._cache["inv_std"]
        N = np.prod(self.normalized_shape)
        if self.affine:
            self.gamma.grad += np.sum(dy * x_centered * inv_std, axis=0)
            self.beta.grad += np.sum(dy, axis=0)
            d_xhat = dy * self.gamma.data
        else:
            d_xhat = dy
        dx = (1./N) * inv_std * (
            N*d_xhat - np.sum(d_xhat, axis=1, keepdims=True)
            - x_centered * (inv_std**2) * np.sum(d_xhat * x_centered, axis=1, keepdims=True)
        )
        return dx

@register_layer("GlobalAvgPool2D")
class GlobalAvgPool2D(Layer):
    def build(self, input_shape, **kwargs):
        N, C, H, W = input_shape
        return (N, C)

    def forward(self, x):
        self._cache["in_shape"] = x.shape
        return np.mean(x, axis=(2,3))

    def backward(self, dy):
        N, C, H, W = self._cache["in_shape"]
        return dy[:, :, None, None] * np.ones((N, C, H, W), dtype=DEFAULT_DTYPE) / (H * W)

@register_layer("ConvTranspose2D")
class ConvTranspose2D(Layer):
    def __init__(self, filters: int, kernel_size: int | Tuple[int,int],
                 stride: int | Tuple[int,int] = 1, padding: str | int | Tuple[int,int] = "valid",
                 use_bias: bool = True, kernel_initializer: str = "he_normal",
                 bias_initializer: str = "zeros", seed: Optional[int] = None):
        super().__init__()
        self.filters = filters
        if isinstance(kernel_size, int):
            self.kH = self.kW = kernel_size
        else:
            self.kH, self.kW = kernel_size
        if isinstance(stride, int):
            self.sH = self.sW = stride
        else:
            self.sH, self.sW = stride
        self.padding = padding
        self.use_bias = use_bias
        self.kernel_initializer = kernel_initializer
        self.bias_initializer = bias_initializer
        self.seed = seed

    def build(self, input_shape, **kwargs):
        N, C, H, W = input_shape
        Wshape = (C, self.filters, self.kH, self.kW)  # note: swapped in/out channels
        W = get_initializer(self.kernel_initializer)(Wshape, seed=self.seed)
        self.W = Parameter(W, name="kernel")
        self._params.append(self.W)
        if self.use_bias:
            b = get_initializer(self.bias_initializer)((self.filters,), seed=self.seed)
            self.b = Parameter(b, name="bias")
            self._params.append(self.b)
        else:
            self.b = None
        out_h = (H - 1) * self.sH + self.kH
        out_w = (W - 1) * self.sW + self.kW
        return (N, self.filters, out_h, out_w)

    def forward(self, x):
        self._cache["x"] = x
        N, C, H, W = x.shape
        out_h = (H - 1) * self.sH + self.kH
        out_w = (W - 1) * self.sW + self.kW
        out = np.zeros((N, self.filters, out_h, out_w), dtype=DEFAULT_DTYPE)
        for n in range(N):
            for c in range(C):
                for f in range(self.filters):
                    out[n, f] += np.kron(x[n, c], np.ones((self.sH, self.sW))) * self.W.data[c, f]
        if self.b is not None:
            out += self.b.data[None, :, None, None]
        return out

    def backward(self, dy):
        x = self._cache["x"]
        N, C, H, W = x.shape
        dW = np.zeros_like(self.W.data)
        dx = np.zeros_like(x)
        for n in range(N):
            for c in range(C):
                for f in range(self.filters):
                    dW[c, f] += np.sum(
                        np.kron(x[n, c], np.ones((self.sH, self.sW))) * dy[n, f]
                    )
                    dx[n, c] += np.convolve(dy[n, f], self.W.data[c, f], mode="valid")
        self.W.grad += dW
        if self.b is not None:
            self.b.grad += np.sum(dy, axis=(0,2,3))
        return dx


@register_layer("Flatten")
class Flatten(Layer):
    def build(self, input_shape, **kwargs):
        self.input_shape = input_shape
        N = input_shape[0]
        out = int(np.prod(input_shape[1:]))
        self.built = True
        return (N, out)

    def forward(self, x):
        self._cache["in_shape"] = x.shape
        return x.reshape(x.shape[0], -1)

    def backward(self, dy):
        return dy.reshape(self._cache["in_shape"])

@register_layer("Dropout")
class Dropout(Layer):
    def __init__(self, rate: float = 0.5, seed: Optional[int] = None):
        super().__init__()
        assert 0.0 <= rate < 1.0
        self.rate = rate
        self.seed = seed

    def forward(self, x):
        if self.training and self.rate > 0.0:
            rng = _get_rng(self.seed)
            mask = rng.random(x.shape) >= self.rate
            scale = 1.0 / (1.0 - self.rate)
            self._cache["mask"] = mask.astype(DEFAULT_DTYPE) * scale
            return x * self._cache["mask"]
        self._cache["mask"] = np.ones_like(x, dtype=DEFAULT_DTYPE)
        return x

    def backward(self, dy):
        return dy * self._cache["mask"]

@register_layer("BatchNorm")
class BatchNorm(Layer):
    # Supports 2D (N, D) or 4D (N, C, H, W) with axis=1 as channel
    def __init__(self, momentum: float = 0.9, eps: float = 1e-5, affine: bool = True):
        super().__init__()
        self.momentum = momentum
        self.eps = eps
        self.affine = affine

    def build(self, input_shape, **kwargs):
        if len(input_shape) == 2:
            C = input_shape[1]
            param_shape = (C,)
        elif len(input_shape) == 4:
            C = input_shape[1]
            param_shape = (C, 1, 1)
        else:
            raise ValueError("BatchNorm expects 2D or 4D input")
        self.running_mean = np.zeros(param_shape, dtype=DEFAULT_DTYPE)
        self.running_var = np.ones(param_shape, dtype=DEFAULT_DTYPE)
        if self.affine:
            self.gamma = Parameter(np.ones(param_shape, dtype=DEFAULT_DTYPE), name="gamma")
            self.beta = Parameter(np.zeros(param_shape, dtype=DEFAULT_DTYPE), name="beta")
            self._params += [self.gamma, self.beta]
        else:
            self.gamma = None
            self.beta = None
        self.axis = 1  # channel
        self.input_dim = len(input_shape)
        self.built = True
        return input_shape

    def forward(self, x):
        if self.input_dim == 2:
            axes = (0,)
            keep = (1,)
        else:
            axes = (0, 2, 3)
            keep = (1, None, None)
        if self.training:
            mean = np.mean(x, axis=axes, keepdims=True)
            var = np.var(x, axis=axes, keepdims=True)
            self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * mean
            self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var
            self._cache["x_centered"] = x - mean
            self._cache["inv_std"] = 1.0 / np.sqrt(var + self.eps)
            x_hat = self._cache["x_centered"] * self._cache["inv_std"]
        else:
            x_hat = (x - self.running_mean) / np.sqrt(self.running_var + self.eps)
        if self.affine:
            out = x_hat * self.gamma.data + self.beta.data
        else:
            out = x_hat
        self._cache["x_hat"] = x_hat
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        x_hat = self._cache["x_hat"]
        if self.affine:
            self.gamma.grad += np.sum(dy * x_hat, axis=(0,2,3), keepdims=True) if x_hat.ndim==4 else np.sum(dy * x_hat, axis=0, keepdims=True)
            self.beta.grad += np.sum(dy, axis=(0,2,3), keepdims=True) if dy.ndim==4 else np.sum(dy, axis=0, keepdims=True)
            d_xhat = dy * self.gamma.data
        else:
            d_xhat = dy

        inv_std = self._cache["inv_std"]
        x_centered = self._cache["x_centered"]
        N = np.prod(dy.shape) / dy.shape[1]  # elements per channel

        dx = (1./N) * inv_std * (N * d_xhat - np.sum(d_xhat, axis=(0,2,3), keepdims=True) if d_xhat.ndim==4 else np.sum(d_xhat, axis=0, keepdims=True) - x_centered * (inv_std**2) * (np.sum(d_xhat * x_centered, axis=(0,2,3), keepdims=True) if d_xhat.ndim==4 else np.sum(d_xhat * x_centered, axis=0, keepdims=True)))
        return dx

# ------------------------------------------------------------
# Losses
# ------------------------------------------------------------

LOSS_REGISTRY: Dict[str, "Loss"] = {}

def register_loss(name: str):
    def deco(cls):
        LOSS_REGISTRY[name] = cls
        return cls
    return deco

class Loss:
    def forward(self, y_pred: np.ndarray, y_true: np.ndarray) -> float:
        raise NotImplementedError
    def backward(self) -> np.ndarray:
        raise NotImplementedError

@register_loss("mse")
class MSELoss(Loss):
    def forward(self, y_pred, y_true):
        self._diff = (y_pred - y_true).astype(DEFAULT_DTYPE)
        return float(np.mean(self._diff ** 2) / 2.0)
    def backward(self):
        N = self._diff.shape[0]
        return self._diff / N

def softmax_stable(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = logits - np.max(logits, axis=axis, keepdims=True)
    exp = np.exp(z)
    return exp / np.sum(exp, axis=axis, keepdims=True)

@register_loss("cross_entropy")
class CrossEntropyLoss(Loss):
    # Expects logits; internally applies softmax; y_true as integer class ids or one-hot
    def __init__(self, axis: int = -1, from_logits: bool = True):
        self.axis = axis
        self.from_logits = from_logits

    def forward(self, y_pred, y_true):
        # y_pred: logits or probabilities
        if self.from_logits:
            probs = softmax_stable(y_pred, axis=self.axis)
        else:
            probs = y_pred / np.clip(np.sum(y_pred, axis=self.axis, keepdims=True), 1e-12, None)
        self._probs = probs.astype(DEFAULT_DTYPE)

        # y_true could be ints or one-hot
        if y_true.ndim == y_pred.ndim:
            # one-hot
            eps = 1e-12
            loss = -np.sum(y_true * np.log(np.clip(probs, eps, 1.0 - eps)), axis=self.axis)
        else:
            # class indices
            idx = y_true.astype(np.int64)
            # gather per-sample prob
            range_idx = np.arange(y_pred.shape[0])
            p = probs[range_idx, idx]
            eps = 1e-12
            loss = -np.log(np.clip(p, eps, 1.0 - eps))
        return float(np.mean(loss))

    def backward(self, y_true=None):
        # If y_true not provided, assume last forward y_true was provided and stored as one-hot? We'll accept y_true in model.backward
        raise NotImplementedError("CrossEntropyLoss.backward requires fused implementation in Model._compute_loss_and_grad")

@register_loss("binary_cross_entropy")
class BinaryCrossEntropy(Loss):
    def forward(self, y_pred, y_true):
        # y_pred as logits OR probabilities? We'll assume probabilities in (0,1)
        eps = 1e-12
        y_pred = np.clip(y_pred, eps, 1 - eps)
        self._y_pred = y_pred
        self._y_true = y_true
        loss = - (y_true*np.log(y_pred) + (1-y_true)*np.log(1-y_pred))
        return float(np.mean(loss))

    def backward(self):
        N = self._y_true.shape[0]
        return (self._y_pred - self._y_true) / (N * np.clip(self._y_pred*(1-self._y_pred), 1e-12, None))

# ------------------------------------------------------------
# Regularization
# ------------------------------------------------------------

def l2_regularization(params: Iterable[Parameter], weight_decay: float) -> Tuple[float, Dict[int, np.ndarray]]:
    reg = 0.0
    grads = {}
    for p in params:
        reg += 0.5 * weight_decay * float(np.sum(p.data**2))
        grads[id(p)] = weight_decay * p.data
    return reg, grads

def l1_regularization(params: Iterable[Parameter], weight_decay: float) -> Tuple[float, Dict[int, np.ndarray]]:
    reg = 0.0
    grads = {}
    for p in params:
        reg += weight_decay * float(np.sum(np.abs(p.data)))
        grads[id(p)] = weight_decay * np.sign(p.data)
    return reg, grads

# ------------------------------------------------------------
# Optimizers
# ------------------------------------------------------------

OPTIMIZER_REGISTRY: Dict[str, "Optimizer"] = {}

def register_optimizer(name: str):
    def deco(cls):
        OPTIMIZER_REGISTRY[name] = cls
        return cls
    return deco

class Optimizer:
    def step(self, params: Iterable[Parameter]):
        raise NotImplementedError
    def zero_grad(self, params: Iterable[Parameter]):
        for p in params:
            p.zero_grad()

@register_optimizer("SGD")
class SGD(Optimizer):
    def __init__(self, lr: float = 1e-2, momentum: float = 0.0, nesterov: bool = False, weight_decay: float = 0.0):
        self.lr = lr
        self.momentum = momentum
        self.nesterov = nesterov
        self.weight_decay = weight_decay
        self._v: Dict[int, np.ndarray] = {}

    def step(self, params: Iterable[Parameter]):
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            if self.momentum != 0.0:
                v = self._v.get(pid, np.zeros_like(g))
                v = self.momentum * v + g
                self._v[pid] = v
                if self.nesterov:
                    update = self.momentum * v + g
                else:
                    update = v
                p.data -= self.lr * update
            else:
                p.data -= self.lr * g

@register_optimizer("Adam")
class Adam(Optimizer):
    def __init__(self, lr: float = 1e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, weight_decay: float = 0.0):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self._m: Dict[int, np.ndarray] = {}
        self._v: Dict[int, np.ndarray] = {}
        self._t = 0

    def step(self, params: Iterable[Parameter]):
        self._t += 1
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            m = self._m.get(pid, np.zeros_like(g))
            v = self._v.get(pid, np.zeros_like(g))
            m = self.beta1 * m + (1 - self.beta1) * g
            v = self.beta2 * v + (1 - self.beta2) * (g * g)
            m_hat = m / (1 - self.beta1 ** self._t)
            v_hat = v / (1 - self.beta2 ** self._t)
            p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
            self._m[pid] = m
            self._v[pid] = v

@register_optimizer("RMSProp")
class RMSProp(Optimizer):
    def __init__(self, lr: float = 1e-3, beta: float = 0.9, eps: float = 1e-8, weight_decay: float = 0.0):
        self.lr = lr
        self.beta = beta
        self.eps = eps
        self.weight_decay = weight_decay
        self._avg_sq: Dict[int, np.ndarray] = {}

    def step(self, params: Iterable[Parameter]):
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            v = self._avg_sq.get(pid, np.zeros_like(g))
            v = self.beta * v + (1 - self.beta) * (g * g)
            self._avg_sq[pid] = v
            p.data -= self.lr * g / (np.sqrt(v) + self.eps)

@register_optimizer("Adagrad")
class Adagrad(Optimizer):
    def __init__(self, lr: float = 1e-2, eps: float = 1e-8, weight_decay: float = 0.0):
        self.lr = lr
        self.eps = eps
        self.weight_decay = weight_decay
        self._G: Dict[int, np.ndarray] = {}

    def step(self, params: Iterable[Parameter]):
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            G = self._G.get(pid, np.zeros_like(g))
            G += g * g
            self._G[pid] = G
            p.data -= self.lr * g / (np.sqrt(G) + self.eps)

@register_optimizer("Adadelta")
class Adadelta(Optimizer):
    def __init__(self, lr: float = 1.0, rho: float = 0.95, eps: float = 1e-6, weight_decay: float = 0.0):
        self.lr = lr
        self.rho = rho
        self.eps = eps
        self.weight_decay = weight_decay
        self._Eg: Dict[int, np.ndarray] = {}
        self._Edx: Dict[int, np.ndarray] = {}

    def step(self, params: Iterable[Parameter]):
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            Eg = self._Eg.get(pid, np.zeros_like(g))
            Edx = self._Edx.get(pid, np.zeros_like(g))
            Eg = self.rho * Eg + (1 - self.rho) * (g * g)
            dx = -np.sqrt((Edx + self.eps) / (Eg + self.eps)) * g
            Edx = self.rho * Edx + (1 - self.rho) * (dx * dx)
            self._Eg[pid] = Eg
            self._Edx[pid] = Edx
            p.data += self.lr * dx

@register_optimizer("Nadam")
class Nadam(Optimizer):
    def __init__(self, lr: float = 1e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, weight_decay: float = 0.0):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self._m: Dict[int, np.ndarray] = {}
        self._v: Dict[int, np.ndarray] = {}
        self._t = 0

    def step(self, params: Iterable[Parameter]):
        self._t += 1
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            m = self._m.get(pid, np.zeros_like(g))
            v = self._v.get(pid, np.zeros_like(g))
            m = self.beta1 * m + (1 - self.beta1) * g
            v = self.beta2 * v + (1 - self.beta2) * (g * g)
            m_hat = m / (1 - self.beta1 ** self._t)
            v_hat = v / (1 - self.beta2 ** self._t)
            m_nesterov = (self.beta1 * m_hat) + ((1 - self.beta1) * g / (1 - self.beta1 ** self._t))
            p.data -= self.lr * m_nesterov / (np.sqrt(v_hat) + self.eps)
            self._m[pid] = m
            self._v[pid] = v

@register_optimizer("Adamax")
class Adamax(Optimizer):
    def __init__(self, lr: float = 2e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, weight_decay: float = 0.0):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self._m: Dict[int, np.ndarray] = {}
        self._u: Dict[int, np.ndarray] = {}
        self._t = 0

    def step(self, params: Iterable[Parameter]):
        self._t += 1
        for p in params:
            g = p.grad
            if self.weight_decay != 0.0:
                g = g + self.weight_decay * p.data
            pid = id(p)
            m = self._m.get(pid, np.zeros_like(g))
            u = self._u.get(pid, np.zeros_like(g))
            m = self.beta1 * m + (1 - self.beta1) * g
            u = np.maximum(self.beta2 * u, np.abs(g))
            m_hat = m / (1 - self.beta1 ** self._t)
            p.data -= (self.lr / (u + self.eps)) * m_hat
            self._m[pid] = m
            self._u[pid] = u


# ------------------------------------------------------------
# Model
# ------------------------------------------------------------

class Model:
    def __init__(self, layers: List[Layer], loss: str = "cross_entropy", loss_kwargs: Optional[Dict[str, Any]] = None,
                 optimizer: str = "Adam", optimizer_kwargs: Optional[Dict[str, Any]] = None,
                 metrics: Optional[List[str]] = None, regularization: Optional[Dict[str, Any]] = None,
                 seed: Optional[int] = None):
        self.layers = layers
        self.loss_name = loss
        self.loss_kwargs = loss_kwargs or {}
        self.loss = self._make_loss(loss, self.loss_kwargs)
        self.optimizer_name = optimizer
        self.optimizer_kwargs = optimizer_kwargs or {}
        self.optimizer = self._make_optimizer(optimizer, self.optimizer_kwargs)
        self.metrics = metrics or []
        self.regularization = regularization or {}
        self.seed = seed

    def _make_loss(self, name, kwargs):
        if name not in LOSS_REGISTRY:
            raise ValueError(f"Unknown loss: {name}")
        return LOSS_REGISTRY[name](**kwargs) if kwargs else LOSS_REGISTRY[name]()

    def _make_optimizer(self, name, kwargs):
        if name not in OPTIMIZER_REGISTRY:
            raise ValueError(f"Unknown optimizer: {name}")
        return OPTIMIZER_REGISTRY[name](**kwargs) if kwargs else OPTIMIZER_REGISTRY[name]()

    def build(self, input_shape: Tuple[int, ...]):
        shape = input_shape
        for layer in self.layers:
            shape = layer.build(shape)
        return shape

    def set_training(self, mode: bool):
        for l in self.layers:
            l.set_training(mode)

    def forward(self, x, training: bool = True):
        self.set_training(training)
        out = x
        for l in self.layers:
            out = l.forward(out)
        return out

    def backward(self, dL):
        grad = dL
        for l in reversed(self.layers):
            grad = l.backward(grad)
        return grad

    def params(self):
        for l in self.layers:
            for p in l.params():
                yield p

    def zero_grads(self):
        for l in self.layers:
            l.zero_grads()

    def _compute_loss_and_grad(self, y_pred, y_true):
        # Supports fused CE+Softmax backward when using CrossEntropyLoss with from_logits=True
        if isinstance(self.loss, CrossEntropyLoss) and self.loss.from_logits:
            # Compute softmax
            probs = softmax_stable(y_pred, axis=-1)
            N = y_pred.shape[0]
            if y_true.ndim == y_pred.ndim:
                # one-hot
                loss = -np.mean(np.sum(y_true * np.log(np.clip(probs, 1e-12, 1.0)), axis=-1))
                dL = (probs - y_true) / N
            else:
                idx = y_true.astype(np.int64)
                loss = -np.mean(np.log(np.clip(probs[np.arange(N), idx], 1e-12, 1.0)))
                dL = probs
                dL[np.arange(N), idx] -= 1
                dL /= N
        else:
            loss = self.loss.forward(y_pred, y_true)
            dL = self.loss.backward()
        # Regularization
        reg_cfg = self.regularization or {}
        reg_loss = 0.0
        reg_grads = {}
        if "l2" in reg_cfg and reg_cfg["l2"] > 0:
            r, g = l2_regularization(self.params(), reg_cfg["l2"])
            reg_loss += r
            reg_grads.update(g)
        if "l1" in reg_cfg and reg_cfg["l1"] > 0:
            r, g = l1_regularization(self.params(), reg_cfg["l1"])
            reg_loss += r
            # accumulate if both present
            for k, v in g.items():
                reg_grads[k] = reg_grads.get(k, 0) + v
        return loss + reg_loss, dL, reg_grads

    def fit(self, x, y, epochs: int = 1, batch_size: int = 32, shuffle: bool = True, verbose: int = 1, callbacks: Optional[List[Any]] = None):
        N = x.shape[0]
        indices = np.arange(N)
        callbacks = callbacks or []

        for cb in callbacks:
            if hasattr(cb, "on_train_begin"):
                cb.on_train_begin(self)

        for epoch in range(1, epochs+1):
            if shuffle:
                np.random.shuffle(indices)
            total_loss = 0.0
            num_batches = int(np.ceil(N / batch_size))

            for b in range(num_batches):
                batch_idx = indices[b*batch_size:(b+1)*batch_size]
                xb = x[batch_idx]
                yb = y[batch_idx]

                # forward
                y_pred = self.forward(xb, training=True)
                # loss + grad
                loss, dL, reg_grads = self._compute_loss_and_grad(y_pred, yb)
                total_loss += loss

                # backward
                self.zero_grads()
                self.backward(dL)
                # add reg grads
                for p in self.params():
                    if id(p) in reg_grads:
                        p.grad += reg_grads[id(p)]
                # step
                self.optimizer.step(self.params())

            avg_loss = total_loss / num_batches
            metric_vals = self._compute_metrics(x, y)
            if verbose:
                msg = f"Epoch {epoch}/{epochs} - loss: {avg_loss:.4f}"
                for k, v in metric_vals.items():
                    msg += f" - {k}: {v:.4f}"
                print(msg)

            for cb in callbacks:
                if hasattr(cb, "on_epoch_end"):
                    cb.on_epoch_end(self, epoch, logs={"loss": avg_loss, **metric_vals})

        for cb in callbacks:
            if hasattr(cb, "on_train_end"):
                cb.on_train_end(self)

    def evaluate(self, x, y, batch_size: int = 128):
        N = x.shape[0]
        num_batches = int(np.ceil(N / batch_size))
        losses = []
        all_metrics = {m: [] for m in self.metrics}
        self.set_training(False)
        for b in range(num_batches):
            xb = x[b*batch_size:(b+1)*batch_size]
            yb = y[b*batch_size:(b+1)*batch_size]
            y_pred = self.forward(xb, training=False)
            loss, _, _ = self._compute_loss_and_grad(y_pred, yb)
            losses.append(loss)
            mv = self._compute_batch_metrics(y_pred, yb)
            for k, v in mv.items():
                all_metrics[k].append(v)
        results = {"loss": float(np.mean(losses))}
        for k, vlist in all_metrics.items():
            if len(vlist) > 0:
                results[k] = float(np.mean(vlist))
        return results

    def predict(self, x, batch_size: int = 128):
        N = x.shape[0]
        num_batches = int(np.ceil(N / batch_size))
        outs = []
        self.set_training(False)
        for b in range(num_batches):
            xb = x[b*batch_size:(b+1)*batch_size]
            outs.append(self.forward(xb, training=False))
        return np.concatenate(outs, axis=0)

    # ---- Metrics ----
    def _compute_metrics(self, x, y):
        y_pred = self.predict(x)
        return self._compute_batch_metrics(y_pred, y)

    def _compute_batch_metrics(self, y_pred, y_true):
        res = {}
        for m in self.metrics:
            if m == "accuracy":
                res[m] = accuracy(y_pred, y_true)
            elif m == "binary_accuracy":
                res[m] = binary_accuracy(y_pred, y_true)
            elif m.startswith("top_"):
                try:
                    k = int(m.split("_")[1])
                    res[m] = top_k_accuracy(y_pred, y_true, k=k)
                except:
                    pass
        return res

    # ---- Serialization ----
    def save_weights(self, path: str):
        arrays = {}
        idx = 0
        for l in self.layers:
            for p in l.params():
                arrays[f"{idx}:{p.name or 'param'}"] = p.data
                idx += 1
        np.savez(path, **arrays)

    def load_weights(self, path: str):
        data = np.load(path)
        idx = 0
        for l in self.layers:
            for p in l.params():
                key = f"{idx}:{p.name or 'param'}"
                if key not in data:
                    raise KeyError(f"Missing weight '{key}' in file")
                p.data[...] = data[key]
                idx += 1

# ------------------------------------------------------------
# Metrics
# ------------------------------------------------------------

def accuracy(y_pred, y_true):
    # y_pred: logits or probabilities
    if y_pred.ndim == 2:
        pred = np.argmax(y_pred, axis=1)
    else:
        pred = (y_pred > 0.5).astype(np.int64).squeeze()
    if y_true.ndim > 1:
        y_true_idx = np.argmax(y_true, axis=1)
    else:
        y_true_idx = y_true.astype(np.int64)
    return float(np.mean(pred == y_true_idx))

def binary_accuracy(y_pred, y_true, threshold: float = 0.5):
    pred = (y_pred >= threshold).astype(np.int64)
    return float(np.mean(pred == y_true))

def top_k_accuracy(y_pred, y_true, k=5):
    idx = np.argpartition(y_pred, -k, axis=1)[:, -k:]
    if y_true.ndim > 1:
        true_idx = np.argmax(y_true, axis=1)
    else:
        true_idx = y_true.astype(np.int64)
    hits = np.any(idx == true_idx[:, None], axis=1)
    return float(np.mean(hits))

# ------------------------------------------------------------
# Builder from config (registry/factory pattern)
# ------------------------------------------------------------

def _maybe_add_activation(layer_list: List[Layer], act_name: Optional[str]):
    if not act_name:
        return
    act_name = act_name.lower()
    if act_name not in ACTIVATION_REGISTRY:
        raise ValueError(f"Unknown activation: {act_name}")
    # Create a new instance (activation layers have __init__ with defaults)
    layer_cls = ACTIVATION_REGISTRY[act_name].__class__ if isinstance(ACTIVATION_REGISTRY[act_name], Layer) else ACTIVATION_REGISTRY[act_name]
    # Handle our registration decoration returning classes; ensure callability
    if isinstance(ACTIVATION_REGISTRY[act_name], type):
        layer_list.append(ACTIVATION_REGISTRY[act_name]())
    else:
        layer_list.append(ACTIVATION_REGISTRY[act_name]())

def build_model_from_config(config: Dict[str, Any]) -> Model:
    layers_cfg = config.get("layers", [])
    layers: List[Layer] = []
    for lc in layers_cfg:
        lc = dict(lc)  # copy
        ltype = lc.pop("type")
        activation = lc.pop("activation", None)
        if ltype not in LAYER_REGISTRY:
            raise ValueError(f"Unknown layer type: {ltype}")
        layer = LAYER_REGISTRY[ltype](**lc)
        layers.append(layer)
        if activation:
            # append activation layer
            act_key = activation.lower()
            if act_key not in ACTIVATION_REGISTRY:
                raise ValueError(f"Unknown activation: {activation}")
            act_cls = ACTIVATION_REGISTRY[act_key]
            # If registry stores class, instantiate; if instance factory, call it
            if isinstance(act_cls, type):
                layers.append(act_cls())
            else:
                layers.append(act_cls())

    loss = config.get("loss", "cross_entropy")
    loss_kwargs = config.get("loss_kwargs", {})
    optimizer_cfg = config.get("optimizer", {"type": "Adam", "lr": 1e-3})
    if isinstance(optimizer_cfg, str):
        optimizer_type = optimizer_cfg
        optimizer_kwargs = {}
    else:
        optimizer_type = optimizer_cfg.get("type", "Adam")
        optimizer_kwargs = {k:v for k,v in optimizer_cfg.items() if k!="type"}
    metrics = config.get("metrics", [])
    regularization = config.get("regularization", {})
    seed = config.get("seed", None)

    model = Model(layers, loss=loss, loss_kwargs=loss_kwargs,
                  optimizer=optimizer_type, optimizer_kwargs=optimizer_kwargs,
                  metrics=metrics, regularization=regularization, seed=seed)

    # Build shapes
    input_shape = tuple(config["input_shape"])
    model.build(input_shape)
    return model

# ------------------------------------------------------------
# Example (runnable) usage
# ------------------------------------------------------------

if __name__ == "__main__":
    # Tiny sanity check: Overfit a toy classification
    rng = np.random.default_rng(42)
    N, D_in, D_hidden, D_out = 64, 20, 32, 5
    x = rng.standard_normal((N, D_in)).astype(DEFAULT_DTYPE)
    # Make labels dependent on first few dims
    logits_true = x[:, :D_out] + 0.1*rng.standard_normal((N, D_out))
    y_idx = np.argmax(logits_true, axis=1).astype(np.int64)

    config = {
        "input_shape": [N, D_in],
        "layers": [
            {"type": "Dense", "units": D_hidden, "activation": "relu"},
            {"type": "Dropout", "rate": 0.1},
            {"type": "Dense", "units": D_out},  # logits
        ],
        "loss": "cross_entropy",
        "optimizer": {"type": "Adam", "lr": 1e-2},
        "metrics": ["accuracy"],
        "regularization": {"l2": 1e-4},
        "seed": 123
    }

    model = build_model_from_config(config)
    model.fit(x, y_idx, epochs=10, batch_size=16, verbose=1)
    res = model.evaluate(x, y_idx)
    print("Eval:", res)