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

from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple

import numpy as np

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


def _get_rng(seed: Optional[int]) -> np.random.Generator:
    """
    Retrieves or creates a NumPy random number generator.

    Args:
        seed: Optional integer seed for reproducibility.

    Returns:
        A NumPy Generator instance.
    """
    return np.random.default_rng(seed) if seed is not None else np.random.default_rng()


DEFAULT_DTYPE = np.float32

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


class Parameter:
    """
    Wrapper for trainable parameters, managing data and gradients.
    """

    def __init__(self, data: np.ndarray, name: Optional[str] = None):
        """
        Initializes the parameter.

        Args:
            data: Initial value of the parameter.
            name: Optional descriptive name for the parameter.
        """
        self.data = data.astype(DEFAULT_DTYPE, copy=False)
        self.grad = np.zeros_like(self.data)
        self.name = name

    def zero_grad(self):
        """Resets the gradient to zero."""
        self.grad[...] = 0


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

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


def register_initializer(name: str):
    """
    Decorator to register an initialization function.

    Args:
        name: Key to use in the initializer registry.
    """

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

    return deco


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


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


@register_initializer("random_normal")
def init_random_normal(shape: Tuple[int, ...], seed: Optional[int] = None):
    """Initializes weights with a standard normal distribution."""
    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):
    """Initializes weights using Glorot (Xavier) Uniform distribution."""
    # 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):
    """Initializes weights using He Normal distribution (best for ReLU)."""
    # 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]:
    """
    Retrieves an initializer function from the registry.

    Args:
        name: Name of the initializer.

    Returns:
        The initialization function.
    """
    if name not in INITIALIZER_REGISTRY:
        raise ValueError(f"Unknown initializer: {name}")
    return INITIALIZER_REGISTRY[name]


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


class Layer:
    """Base class for all layers."""

    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, ...]:
        """
        Builds layer structure given input shape.

        Returns:
            The output shape of the layer.
        """
        self.built = True
        return input_shape

    def forward(self, x: np.ndarray) -> np.ndarray:
        """Computes the forward pass of the layer."""
        raise NotImplementedError

    def backward(self, dy: np.ndarray) -> np.ndarray:
        """Computes the backward pass (gradients) of the layer."""
        raise NotImplementedError

    def params(self) -> Iterable[Parameter]:
        """Returns list of trainable parameters in this layer."""
        return self._params

    def zero_grads(self):
        """Resets gradients of all parameters in this layer."""
        for p in self._params:
            p.zero_grad()

    def set_training(self, mode: bool):
        """Toggles training mode for layers that behave differently during training vs inference."""
        self.training = mode

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


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


def register_layer(name: str):
    """Decorator to register a neural network layer."""

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

    return deco


def register_activation(name: str):
    """Decorator to register an activation function."""

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

    return deco


# ---- Activation implementations ----
@register_activation("relu")
@register_layer("ReLU")
class ReLU(Layer):
    """Rectified Linear Unit activation."""

    def forward(self, x):
        """
        Applies the ReLU activation function.

        Args:
            x (np.ndarray): Input array.

        Returns:
            np.ndarray: Output after applying ReLU.
        """
        self._cache["mask"] = x > 0
        return np.where(self._cache["mask"], x, 0).astype(DEFAULT_DTYPE)

    def backward(self, dy):
        """
        Computes the gradient of the ReLU function.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        return dy * self._cache["mask"]


@register_activation("sigmoid")
@register_layer("Sigmoid")
class Sigmoid(Layer):
    """Sigmoid activation function."""

    def forward(self, x):
        """
        Applies the Sigmoid activation function.

        Args:
            x (np.ndarray): Input array.

        Returns:
            np.ndarray: Output after applying Sigmoid.
        """
        out = 1.0 / (1.0 + np.exp(-x))
        self._cache["out"] = out
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        """
        Computes the gradient of the Sigmoid function.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        out = self._cache["out"]
        return dy * out * (1.0 - out)


@register_activation("tanh")
@register_layer("Tanh")
class Tanh(Layer):
    """Hyperbolic tangent activation function."""

    def forward(self, x):
        """
        Applies the Hyperbolic Tangent activation function.

        Args:
            x (np.ndarray): Input array.

        Returns:
            np.ndarray: Output after applying Tanh.
        """
        out = np.tanh(x)
        self._cache["out"] = out
        return out.astype(DEFAULT_DTYPE)

    def backward(self, dy):
        """
        Computes the gradient of the Tanh function.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        out = self._cache["out"]
        return dy * (1.0 - out**2)


@register_activation("leaky_relu")
@register_layer("LeakyReLU")
class LeakyReLU(Layer):
    """Leaky ReLU activation function."""

    def __init__(self, alpha: float = 0.01):
        """
        Initializes the LeakyReLU layer.

        Args:
            alpha (float): Scaling factor for negative input values.
        """
        super().__init__()
        self.alpha = alpha

    def forward(self, x):
        """
        Applies the Leaky ReLU activation function.

        Args:
            x (np.ndarray): Input array.

        Returns:
            np.ndarray: Output after applying Leaky ReLU.
        """
        self._cache["x"] = x
        return np.where(x > 0, x, self.alpha * x).astype(DEFAULT_DTYPE)

    def backward(self, dy):
        """
        Computes the gradient of the Leaky ReLU function.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        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):
    """Softmax activation function."""

    def __init__(self, axis: int = -1):
        """
        Initializes the Softmax layer.

        Args:
            axis (int): The axis along which to compute the softmax.
        """
        super().__init__()
        self.axis = axis

    def forward(self, x):
        """
        Applies the Softmax activation function (numerically stable).

        Args:
            x (np.ndarray): Input array.

        Returns:
            np.ndarray: Output after applying Softmax.
        """
        # 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):
        """
        Computes the gradient of the Softmax function.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        # 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):
    """Fully connected layer."""

    def __init__(
        self,
        units: int,
        use_bias: bool = True,
        kernel_initializer: str = "glorot_uniform",
        bias_initializer: str = "zeros",
        seed: Optional[int] = None,
    ):
        """
        Initializes the Dense layer.

        Args:
            units (int): Number of output units.
            use_bias (bool): Whether to include a bias vector.
            kernel_initializer (str): Initializer for the kernel weights.
            bias_initializer (str): Initializer for the bias vector.
            seed (int, optional): Random seed for initialization.
        """
        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, ...]:
        """
        Builds the layer parameters based on input shape.

        Args:
            input_shape (tuple): The shape of the input tensor.

        Returns:
            tuple: The output shape of the layer.
        """
        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):
        """
        Performs the forward pass of the Dense layer.

        Args:
            x (np.ndarray): Input tensor.

        Returns:
            np.ndarray: Transformed output tensor.
        """
        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):
        """
        Computes the gradient of the Dense layer.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        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):
    """Calculates output dimension for a spatial operation."""
    # 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):
    """Calculates padding required for 'same' convolution/pooling."""
    # 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
):
    """Transforms image input to column format for efficient convolution."""
    # 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
):
    """Transforms column data back into original image shape."""
    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):
    """2D Convolution 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,
    ):
        """
        Initializes the Conv2D layer.

        Args:
            filters (int): Number of filters.
            kernel_size (int or tuple): Size of the convolution kernel.
            stride (int or tuple): Stride of the convolution.
            padding (str, int, or tuple): Padding strategy or size.
            dilation (int or tuple): Dilation rate.
            use_bias (bool): Whether to include bias.
            kernel_initializer (str): Initializer for kernel weights.
            bias_initializer (str): Initializer for bias weights.
            seed (int, optional): Random seed.
        """
        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):
        """Internal helper to calculate padding."""
        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, ...]:
        """
        Builds the layer parameters and computes output shape.

        Args:
            input_shape (tuple): The input shape (N, C, H, W).

        Returns:
            tuple: The output shape (N, F, out_h, out_w).
        """
        # 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):
        """
        Performs the forward pass of the Conv2D layer.

        Args:
            x (np.ndarray): Input tensor.

        Returns:
            np.ndarray: Transformed output tensor.
        """
        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):
        """
        Computes the gradient of the Conv2D layer.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        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):
    """2D Max Pooling 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",
    ):
        """
        Initializes the MaxPool2D layer.

        Args:
            pool_size (int or tuple): Window size for pooling.
            stride (int or tuple, optional): Stride of pooling.
            padding (str, int, or tuple): Padding strategy.
        """
        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):
        """Internal helper to calculate padding."""
        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):
        """
        Builds the layer and computes output shape.

        Args:
            input_shape (tuple): Input shape.

        Returns:
            tuple: Output shape.
        """
        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):
        """
        Performs the forward pass of the MaxPool2D layer.

        Args:
            x (np.ndarray): Input tensor.

        Returns:
            np.ndarray: Transformed output tensor.
        """
        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):
        """
        Computes the gradient of the MaxPool2D layer.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to the input.
        """
        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):
    """2D Average Pooling 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",
    ):
        """
        Initializes the AvgPool2D layer.

        Args:
            pool_size (int or tuple): The size of the pooling window.
            stride (int or tuple, optional): The stride of the pooling operation.
            padding (str, int, or tuple): Padding configuration.
        """
        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):
        """Internal helper to calculate padding."""
        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):
        """Builds the layer output shape."""
        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):
        """
        Performs the forward pass of the AvgPool2D layer.

        Args:
            x (np.ndarray): Input tensor.

        Returns:
            np.ndarray: Pooled output tensor.
        """
        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):
        """
        Performs the backward pass of the AvgPool2D layer.

        Args:
            dy (np.ndarray): Upstream gradient.

        Returns:
            np.ndarray: Gradient with respect to input.
        """
        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):
    """Layer to change shape of data."""

    def __init__(self, target_shape: Tuple[int, ...]):
        """
        Initializes the Reshape layer.

        Args:
            target_shape (tuple): The target shape excluding the batch dimension.
        """
        super().__init__()
        self.target_shape = tuple(target_shape)
        self.input_shape = None

    def build(self, input_shape, **kwargs):
        """Builds the layer output shape."""
        self.input_shape = input_shape[1:]
        return (input_shape[0],) + self.target_shape

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

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


@register_layer("UpSample2D")
class UpSample2D(Layer):
    """2D Upsampling layer (nearest neighbor)."""

    def __init__(self, scale: int = 2):
        """
        Initializes the UpSample2D layer.

        Args:
            scale (int): The scaling factor.
        """
        super().__init__()
        self.scale = scale

    def build(self, input_shape, **kwargs):
        """Builds the layer output shape."""
        N, C, H, W = input_shape
        return (N, C, H * self.scale, W * self.scale)

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

    def backward(self, dy):
        """Performs the backward pass for upsampling."""
        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):
    """Layer Normalization implementation."""

    def __init__(self, eps: float = 1e-5, affine: bool = True):
        """
        Initializes the LayerNorm layer.

        Args:
            eps (float): Small value for numerical stability.
            affine (bool): Whether to use learnable affine parameters.
        """
        super().__init__()
        self.eps = eps
        self.affine = affine

    def build(self, input_shape, **kwargs):
        """Builds parameters for normalization."""
        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):
        """Performs Layer Normalization."""
        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):
        """Computes gradients for Layer Normalization."""
        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.0 / 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):
    """Global Average Pooling layer."""

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        """
        Builds the layer by defining output dimensions.

        Args:
            input_shape (tuple): Input shape (N, C, H, W).

        Returns:
            tuple: Output shape (N, C).
        """
        N, C, H, W = input_shape
        return (N, C)

    def forward(self, x: np.ndarray) -> np.ndarray:
        """
        Performs global average pooling over the spatial dimensions.

        Args:
            x (np.ndarray): Input tensor (N, C, H, W).

        Returns:
            np.ndarray: Pooled output (N, C).
        """
        self._cache["in_shape"] = x.shape
        return np.mean(x, axis=(2, 3))

    def backward(self, dy: np.ndarray) -> np.ndarray:
        """
        Computes the gradient of the loss with respect to input.

        Args:
            dy (np.ndarray): Gradient of loss w.r.t. layer output.

        Returns:
            np.ndarray: Gradient of loss w.r.t. layer input.
        """
        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):
    """2D Transposed Convolution 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,
    ):
        """
        Initializes the ConvTranspose2D layer.

        Args:
            filters (int): Number of output filters.
            kernel_size (int or tuple): The size of the kernel.
            stride (int or tuple): The stride of the operation.
            padding (str, int, or tuple): Padding configuration.
            use_bias (bool): Whether to use bias.
            kernel_initializer (str): Initializer for kernel weights.
            bias_initializer (str): Initializer for bias weights.
            seed (int, optional): Random seed.
        """
        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: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        """Initializes kernels and determines output shape."""
        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: np.ndarray) -> np.ndarray:
        """Performs forward pass of transposed convolution."""
        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: np.ndarray) -> np.ndarray:
        """Performs backward pass of transposed convolution."""
        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):
    """Flattens input tensors to 2D (batch, features)."""

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        """Calculates flattened output size."""
        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: np.ndarray) -> np.ndarray:
        """Reshapes input to flat format."""
        self._cache["in_shape"] = x.shape
        return x.reshape(x.shape[0], -1)

    def backward(self, dy: np.ndarray) -> np.ndarray:
        """Restores original shape of flattened input."""
        return dy.reshape(self._cache["in_shape"])


@register_layer("Dropout")
class Dropout(Layer):
    """Applies dropout regularization."""

    def __init__(self, rate: float = 0.5, seed: Optional[int] = None):
        """
        Initializes the Dropout layer.

        Args:
            rate (float): Probability of dropping units.
            seed (int, optional): Random seed.
        """
        super().__init__()
        assert 0.0 <= rate < 1.0
        self.rate = rate
        self.seed = seed

    def forward(self, x: np.ndarray) -> np.ndarray:
        """Applies random masking during training."""
        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: np.ndarray) -> np.ndarray:
        """Propagates gradients through the dropout mask."""
        return dy * self._cache["mask"]


@register_layer("BatchNorm")
class BatchNorm(Layer):
    """Batch Normalization 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):
        """
        Initializes the BatchNorm layer.

        Args:
            momentum (float): Momentum for moving averages.
            eps (float): Stability constant.
            affine (bool): Whether to use learnable affine parameters.
        """
        super().__init__()
        self.momentum = momentum
        self.eps = eps
        self.affine = affine

    def build(self, input_shape: Tuple[int, ...], **kwargs) -> Tuple[int, ...]:
        """Builds running stats and optional learnable parameters."""
        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
        self.input_dim = len(input_shape)
        self.built = True
        return input_shape

    def forward(self, x: np.ndarray) -> np.ndarray:
        """Normalizes input and updates running statistics."""
        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: np.ndarray) -> np.ndarray:
        """Calculates gradients for input, gamma, and beta."""
        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.0 / 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):
    """
    Decorator to register a loss class in the LOSS_REGISTRY.

    Args:
        name (str): The name to associate with the loss class.
    """

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

    return deco


class Loss:
    """Base class for all loss functions."""

    def forward(self, y_pred: np.ndarray, y_true: np.ndarray) -> float:
        """
        Computes loss value.

        Args:
            y_pred (np.ndarray): Predicted output.
            y_true (np.ndarray): Target output.

        Returns:
            float: Computed scalar loss value.
        """
        raise NotImplementedError

    def backward(self) -> np.ndarray:
        """
        Computes gradient of the loss.

        Returns:
            np.ndarray: Gradient of the loss with respect to predictions.
        """
        raise NotImplementedError


@register_loss("mse")
class MSELoss(Loss):
    """Mean Squared Error loss function."""

    def forward(self, y_pred, y_true):
        """
        Computes Mean Squared Error loss.

        Args:
            y_pred (np.ndarray): Predicted values.
            y_true (np.ndarray): Target values.

        Returns:
            float: MSE loss.
        """
        self._diff = (y_pred - y_true).astype(DEFAULT_DTYPE)
        return float(np.mean(self._diff**2) / 2.0)

    def backward(self):
        """
        Computes gradient of the MSE loss.

        Returns:
            np.ndarray: Gradient array.
        """
        N = self._diff.shape[0]
        return self._diff / N


def softmax_stable(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    """
    Numerically stable softmax implementation.

    Args:
        logits (np.ndarray): Input logits.
        axis (int): Axis along which to compute softmax.

    Returns:
        np.ndarray: Softmax probabilities.
    """
    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):
    """Cross Entropy loss function, typically for multi-class classification."""

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

        Args:
            axis (int): The axis along which to apply softmax/sum.
            from_logits (bool): Whether input contains raw logits or probabilities.
        """
        self.axis = axis
        self.from_logits = from_logits

    def forward(self, y_pred, y_true):
        """
        Computes Cross Entropy loss.

        Args:
            y_pred (np.ndarray): Logits or probabilities.
            y_true (np.ndarray): Ground truth labels (class indices or one-hot).

        Returns:
            float: Mean cross entropy loss.
        """
        # 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):
        """
        Placeholder for backward pass; requires fused implementation.

        Raises:
            NotImplementedError: Always, as gradient is handled by model.
        """
        # 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):
    """Binary Cross Entropy loss function."""

    def forward(self, y_pred, y_true):
        """
        Computes Binary Cross Entropy loss.

        Args:
            y_pred (np.ndarray): Predicted probabilities (0 to 1).
            y_true (np.ndarray): Target labels (0 or 1).

        Returns:
            float: Mean binary cross entropy loss.
        """
        # 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):
        """
        Computes gradient of the Binary Cross Entropy loss.

        Returns:
            np.ndarray: Gradient array.
        """
        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]]:
    """
    Calculates L2 penalty and its gradients.

    Args:
        params (Iterable[Parameter]): Parameters to regularize.
        weight_decay (float): Strength of regularization.

    Returns:
        Tuple[float, Dict[int, np.ndarray]]: (Penalty value, Dictionary of gradients).
    """
    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]]:
    """
    Calculates L1 penalty and its gradients.

    Args:
        params (Iterable[Parameter]): Parameters to regularize.
        weight_decay (float): Strength of regularization.

    Returns:
        Tuple[float, Dict[int, np.ndarray]]: (Penalty value, Dictionary of gradients).
    """
    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):
    """
    Decorator to register an optimizer class in the OPTIMIZER_REGISTRY.

    Args:
        name (str): Name to register the optimizer under.
    """

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

    return deco


class Optimizer:
    """Base class for all optimizers."""

    def step(self, params: Iterable[Parameter]):
        """
        Performs an optimization step.

        Args:
            params (Iterable[Parameter]): Parameters to update.
        """
        raise NotImplementedError

    def zero_grad(self, params: Iterable[Parameter]):
        """
        Zeros out gradients for all parameters.

        Args:
            params (Iterable[Parameter]): Parameters to reset.
        """
        for p in params:
            p.zero_grad()


@register_optimizer("SGD")
class SGD(Optimizer):
    """Stochastic Gradient Descent with optional momentum and Nesterov."""

    def __init__(
        self,
        lr: float = 1e-2,
        momentum: float = 0.0,
        nesterov: bool = False,
        weight_decay: float = 0.0,
    ):
        """
        Initializes SGD optimizer.
        """
        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]):
        """
        Performs a single SGD optimization step.
        """
        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
                update = (self.momentum * v + g) if self.nesterov else v
                p.data -= self.lr * update
            else:
                p.data -= self.lr * g


@register_optimizer("Adam")
class Adam(Optimizer):
    """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,
    ):
        """Initializes Adam optimizer."""
        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]):
        """Performs a single Adam optimization step."""
        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):
    """RMSProp optimizer."""

    def __init__(
        self,
        lr: float = 1e-3,
        beta: float = 0.9,
        eps: float = 1e-8,
        weight_decay: float = 0.0,
    ):
        """Initializes RMSProp optimizer."""
        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]):
        """Performs a single RMSProp optimization step."""
        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):
    """Adagrad optimizer."""

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

    def step(self, params: Iterable[Parameter]):
        """Performs a single Adagrad optimization step."""
        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):
    """Adadelta optimizer."""

    def __init__(
        self,
        lr: float = 1.0,
        rho: float = 0.95,
        eps: float = 1e-6,
        weight_decay: float = 0.0,
    ):
        """Initializes Adadelta optimizer."""
        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]):
        """Performs a single Adadelta optimization step."""
        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):
    """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,
    ):
        """Initializes Nadam optimizer."""
        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]):
        """Performs a single Nadam optimization step."""
        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):
    """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,
    ):
        """Initializes Adamax optimizer."""
        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]):
        """Performs a single Adamax optimization step."""
        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:
    """Main model class for composing layers and managing training."""

    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,
    ):
        """Initializes Model with layers, loss, optimizer, and configuration."""
        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):
        """Factory method to create a loss object."""
        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):
        """Factory method to create an optimizer object."""
        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, ...]):
        """Builds layer shapes sequentially."""
        shape = input_shape
        for layer in self.layers:
            shape = layer.build(shape)
        return shape

    def set_training(self, mode: bool):
        """Sets the training mode for all layers."""
        for l in self.layers:
            l.set_training(mode)

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

    def backward(self, dL):
        """Propagates loss gradients backward through all layers."""
        grad = dL
        for l in reversed(self.layers):
            grad = l.backward(grad)
        return grad

    def params(self):
        """Yields all trainable parameters."""
        for l in self.layers:
            for p in l.params():
                yield p

    def zero_grads(self):
        """Zeros gradients for all layers."""
        for l in self.layers:
            l.zero_grads()

    def _compute_loss_and_grad(self, y_pred, y_true):
        """Computes loss and backpropagation gradients, including regularization."""
        # 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,
    ):
        """Trains the model."""
        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):
        """Evaluates model performance."""
        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):
        """Generates predictions for input data."""
        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):
        """Computes metrics over the entire dataset."""
        y_pred = self.predict(x)
        return self._compute_batch_metrics(y_pred, y)

    def _compute_batch_metrics(self, y_pred, y_true):
        """Computes metrics for a batch."""
        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):
        """Saves model weights to a file."""
        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):
        """Loads model weights from a file."""
        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):
    """Calculates classification accuracy."""
    # 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):
    """Calculates binary classification accuracy."""
    pred = (y_pred >= threshold).astype(np.int64)
    return float(np.mean(pred == y_true))


def top_k_accuracy(y_pred, y_true, k=5):
    """Calculates Top-K accuracy."""
    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]):
    """Helper to add an activation layer to the model architecture list."""
    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:
    """Builds a Model instance from a dictionary configuration."""
    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