from __future__ import division, print_function, absolute_import
import math,random,warnings

class DOtensor(object):
    """DOtensor is a function for 1D grad (after all, it's called Dearning Of tensor, abbreviated as DOtensor).

    with the following functions:
    - backward, automatic gradient propagation
    - zero_grad, to reset the gradient
    - trace, to trace code"""
    __slots__ = ("data", "shape", "grad", "requires_grad", "_backward", "_prev")
    _global_trace_enabled = False
    _global_trace_log = []
    def __init__(self, data, requires_grad=False):
        if isinstance(data, DOtensor): data, shape = data.data, data.shape
        else: data, shape = DOtensor._flatten(data)
        self.data = data
        self.shape = shape
        self.requires_grad = requires_grad
        self.grad = [0.0] * len(data) if requires_grad else None
        self._prev = ()
        self._backward = None

    @classmethod
    def trace(cls, enabled=True):
        """trace is a function for tracing tensor code (you probably already know this, as the name suggests). This function can be turned on or off as desired.

        example:
        >>> with DOtensor.trace(enabled=True) as log:
        >>>     a = DOtensor([1.0, 2.0], requires_grad=True)
        >>>     b = DOtensor([3.0, 4.0], requires_grad=True)
        >>>     c = a * b
        >>>     d = c.DTS()
        >>>     d.backward()
        >>> for entry in log: print(entry)"""
        class _TraceContext:
            def __enter__(self_inner):
                cls._prev_trace_state = cls._global_trace_enabled
                cls._global_trace_enabled = enabled
                if enabled: cls._global_trace_log = []
                return cls._global_trace_log

            def __exit__(self_inner, exc_type, exc, tb): cls._global_trace_enabled = cls._prev_trace_state
        return _TraceContext()

    def __add__(self, other):
        other = other if isinstance(other, DOtensor) else DOtensor([other])
        if len(self.data) != len(other.data):
            raise ValueError("DOtensor: shape mismatch {} vs {}".format(len(self.data), len(other.data)))
        def backward(self):
            if not self.requires_grad: return
        out = DOtensor([a+b for a,b in zip(self.data, other.data)], requires_grad=self.requires_grad or other.requires_grad)
        if out.requires_grad:
            def _backward():
                if out.grad is None: return
                if self.requires_grad:
                    for i in range(len(self.grad)): self.grad[i] += out.grad[i]
                if other.requires_grad:
                    for i in range(len(other.grad)): other.grad[i] += out.grad[i]
            out._backward = _backward
            out._prev = (self, other)
        if DOtensor._global_trace_enabled:DOtensor._global_trace_log.append(("add", repr(self), repr(other), repr(out)))
        return out
    
    def __mul__(self, other):
        other = other if isinstance(other, DOtensor) else DOtensor([other])
        if len(self.data) != len(other.data):
            raise ValueError("DOtensor: shape mismatch {} vs {}".format(len(self.data), len(other.data)))
        def backward(self):
            if not self.requires_grad: return
        out = DOtensor([a*b for a,b in zip(self.data, other.data)], requires_grad=self.requires_grad or other.requires_grad)
        if out.requires_grad:
            def _backward():
                if out.grad is None: return
                if self.requires_grad:
                    for i in range(len(self.grad)): self.grad[i] += other.data[i] * out.grad[i]
                if other.requires_grad:
                    for i in range(len(other.grad)): other.grad[i] += self.data[i] * out.grad[i]
            out._backward = _backward
            out._prev = (self, other)
        if DOtensor._global_trace_enabled:DOtensor._global_trace_log.append(("mul", repr(self), repr(other), repr(out)))
        return out
    
    def _flatten(data):
        if isinstance(data, (int, float)): return [float(data)], ()
        if isinstance(data, list):
            if not data: return [], (0,)
            if isinstance(data[0], list):
                rows = len(data)
                cols = len(data[0])
                flat = [float(x) for row in data for x in row]
                return flat, (rows, cols)
            else: return [float(x) for x in data], (len(data),)
        raise TypeError("Invalid data type")

    def __sub__(self, other):
        other = other if isinstance(other, DOtensor) else DOtensor([other])
        return self + (-other)

    def __neg__(self):
        out = DOtensor([-x for x in self.data], self.requires_grad)
        if self.requires_grad:
            def _backward():
                for i in range(len(self.data)): self.grad[i] += -out.grad[0]
            out._backward = _backward
            out._prev = (self,)
        return out
    
    def DTS(self):
        """DTS (Dearning Tensor Sum) is a function that sums all tensor elements into a single scalar.

        Example:
        >>> from dearning import DOtensor
        >>> x = DOtensor([2.0, 3.0], requires_grad=True)
        >>> y = DOtensor([4.0, 5.0], requires_grad=True)
        >>> z = x * y + x
        >>> s = z.DTS()"""
        out = DOtensor([sum(self.data)], self.requires_grad)
        if self.requires_grad:
            def _backward():
                for i in range(len(self.data)): self.grad[i] += out.grad[0]
            out._backward = _backward
            out._prev = (self,)
        return out

    def backward(self):
        """Backward is a function that performs backward propagation on a computational graph.

        Example:
        >>> from dearning import DOtensor
        >>> x = DOtensor([2.0, 3.0], requires_grad=True)
        >>> y = DOtensor([4.0, 5.0], requires_grad=True)
        >>> z = x * y + x
        >>> s = z.DTS()
        >>> s.backward()
        Must be called on a scalar tensor (the final result)."""
        topo = []
        visited = set()
        def build(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev: build(child)
                topo.append(v)
        build(self)
        self.grad = [1.0]
        for node in reversed(topo):
            if node._backward: node._backward()
        if DOtensor._global_trace_enabled: DOtensor._global_trace_log.append(("backward", repr(self)))

    def zero_grad(self):
        """zero_grad is a function to reset grad.

        example:
        >>> from dearning import DOtensor
        >>> x = DOtensor([2.0, 3.0], requires_grad=True)
        >>> y = DOtensor([4.0, 5.0], requires_grad=True)
        >>> x.zero_grad()
        >>> y.zero_grad()
        >>> print("x.grad:", x.grad)
        >>> print("y.grad:", y.grad)"""
        visited = set()
        def reset(v):
            if v not in visited:
                visited.add(v)
                if v.requires_grad: v.grad = [0.0] * len(v.data)
                for p in v._prev: reset(p)
        reset(self)

    def __hash__(self): return id(self)
    def __repr__(self): return "DOtensor(data={}, grad={})".format(self.data,self.grad)

class Sequence:
    class RNN:
        """RNN is Recurrent Neural Network.

        example:
        >>> from dearning import Sequence
        >>> rnn = Sequence.RNN(input_dim=3, hidden_dim=5, output_dim=2, return_sequences=False)
        """
        def __init__(self, input_dim, hidden_dim, output_dim=None, return_sequences=False):
            self.input_dim = int(input_dim)
            self.hidden_dim = int(hidden_dim)
            self.output_dim = int(output_dim) if output_dim is not None else None
            self.return_sequences = bool(return_sequences)
            self.Wxh = [[random.gauss(0, 1) * math.sqrt(2.0 / self.input_dim) for _ in range(self.hidden_dim)] for _ in range(self.input_dim)]
            self.Whh = [[random.gauss(0, 1) * math.sqrt(2.0 / self.hidden_dim) for _ in range(self.hidden_dim)] for _ in range(self.hidden_dim)]
            self.bh = [0.0] * self.hidden_dim
            if self.output_dim is not None:
                self.Why = [[random.gauss(0, 1) * math.sqrt(2.0 / self.hidden_dim) for _ in range(self.output_dim)] for _ in range(self.hidden_dim)]
                self.by = [0.0] * self.output_dim
        def __call__(self, x, training=True): return self.forward(x, training)
        @staticmethod
        def _is_3d(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], list)
        @staticmethod
        def _zeros(b, n): return [[0.0] * n for _ in range(b)]
        def _to_batch(self, X):
            if not X: return [], True
            if Sequence.RNN._is_3d(X): return X, False
            return [X], True

        def forward(self, X, training=True):
            Xb, single = self._to_batch(X)
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            z = Sequence.RNN._zeros
            self._cache = {"X": Xb, "Ts": []}
            hs = z(B, H)
            outs = []
            for t in range(T):
                x_t = [Xb[b][t] for b in range(B)]
                h_t = [[math.tanh(self.bh[j] + sum(x_t[b][i] * self.Wxh[i][j] for i in range(self.input_dim)) + sum(hs[b][k] * self.Whh[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                hs = h_t
                self._cache["Ts"].append(hs)
                outs.append([[self.by[j] + sum(hs[b][k] * self.Why[k][j] for k in range(H)) for j in range(self.output_dim)] for b in range(B)] if self.output_dim is not None else [row[:] for row in hs])
            if not self.return_sequences:
                last = outs[-1] if outs else ([] if not single else [])
                return last[0] if single else last
            if single: return [o[0] for o in outs]
            return outs

        def backward(self, grad_output):
            Xb = self._cache["X"]
            Ts = self._cache["Ts"]
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            single = not Sequence.RNN._is_3d(Xb)
            z = Sequence.RNN._zeros
            if not self.return_sequences: # normalize grad_output to sequence form
                g_last = grad_output if isinstance(grad_output[0], list) else [grad_output]
                G = [z(B, self.output_dim or H) for _ in range(T)]
                G[-1] = g_last
            else:
                G = grad_output
                if single: G = [[g] for g in G]
            self.grad_Wxh = [[0.0] * H for _ in range(self.input_dim)]
            self.grad_Whh = [[0.0] * H for _ in range(H)]
            self.grad_bh = [0.0] * H
            if self.output_dim is not None:
                self.grad_Why = [[0.0] * self.output_dim for _ in range(H)]
                self.grad_by = [0.0] * self.output_dim
            dX = [[[0.0] * self.input_dim for _ in range(T)] for _ in range(B)]
            dh_next = z(B, H)
            for t in reversed(range(T)):
                h_t = Ts[t]
                h_prev = Ts[t - 1] if t > 0 else z(B, H)
                if self.output_dim is not None: # output grads
                    dy = G[t]
                    dh = z(B, H)
                    for b in range(B):
                        for j in range(self.output_dim):
                            self.grad_by[j] += dy[b][j]
                            for k in range(H): self.grad_Why[k][j] += h_t[b][k] * dy[b][j]
                        for k in range(H):
                            s = 0.0
                            for j in range(self.output_dim): s += dy[b][j] * self.Why[k][j]
                            dh[b][k] = s
                else: dh = G[t]
                for b in range(B):
                    for k in range(H): dh[b][k] += dh_next[b][k]
                # tanh backprop
                dh_raw = z(B, H)
                for b in range(B):
                    for k in range(H):
                        dh_raw[b][k] = (1.0 - h_t[b][k] * h_t[b][k]) * dh[b][k]
                        self.grad_bh[k] += dh_raw[b][k]
                for b in range(B):
                    for i in range(self.input_dim):
                        for k in range(H): self.grad_Wxh[i][k] += Xb[b][t][i] * dh_raw[b][k]
                    for k in range(H):
                        for j in range(H): self.grad_Whh[j][k] += h_prev[b][j] * dh_raw[b][k]
                dh_next = z(B, H)
                for b in range(B):
                    for i in range(self.input_dim):
                        s = 0.0
                        for k in range(H): s += dh_raw[b][k] * self.Wxh[i][k]
                        dX[b][t][i] = s
                    for j in range(H):
                        s = 0.0
                        for k in range(H): s += dh_raw[b][k] * self.Whh[j][k]
                        dh_next[b][j] = s
            if single: return dX[0]
            return dX

        def update(self, lr):
            for i in range(self.input_dim):
                for j in range(self.hidden_dim): self.Wxh[i][j] -= lr * self.grad_Wxh[i][j]
            for i in range(self.hidden_dim):
                for j in range(self.hidden_dim): self.Whh[i][j] -= lr * self.grad_Whh[i][j]
            for j in range(self.hidden_dim): self.bh[j] -= lr * self.grad_bh[j]
            if self.output_dim is not None:
                for i in range(self.hidden_dim):
                    for j in range(self.output_dim): self.Why[i][j] -= lr * self.grad_Why[i][j]
                for j in range(self.output_dim): self.by[j] -= lr * self.grad_by[j]

    class GRU:
        """GRU is Gated Recurrent Unit.
        
        example:
        >>> from dearning import Sequence
        >>> gru = Sequence.GRU(input_dim=3, hidden_dim=5, output_dim=2, return_sequences=False)"""
        def __init__(self, input_dim, hidden_dim, output_dim=None, return_sequences=False):
            self.input_dim = int(input_dim)
            self.hidden_dim = int(hidden_dim)
            self.output_dim = int(output_dim) if output_dim is not None else None
            self.return_sequences = bool(return_sequences)
            H, D = self.hidden_dim, self.input_dim
            self.Wz = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Uz = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bz = [0.0] * H
            self.Wr = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Ur = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.br = [0.0] * H
            self.Wh = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Uh = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bh = [0.0] * H
            if self.output_dim is not None:
                self.Why = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(self.output_dim)] for _ in range(H)]
                self.by = [0.0] * self.output_dim
        def __call__(self, x, training=True): return self.forward(x, training)
        @staticmethod
        def _sigmoid(x): return 1.0 / (1.0 + math.exp(-x))
        @staticmethod
        def _is_3d(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], list)
        @staticmethod
        def _zeros(b, n): return [[0.0] * n for _ in range(b)]
        def _to_batch(self, X):
            if not X: return [], True
            if Sequence.GRU._is_3d(X): return X, False
            return [X], True

        def forward(self, X, training=True):
            Xb, single = self._to_batch(X)
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            z = Sequence.GRU._zeros
            hs = z(B, H)
            self._cache = {"X": Xb, "Ts": []}
            outs = []
            for t in range(T):
                x_t = [Xb[b][t] for b in range(B)]
                z_t = [[Sequence.GRU._sigmoid(self.bz[j] + sum(x_t[b][i] * self.Wz[i][j] for i in range(self.input_dim)) + sum(hs[b][k] * self.Uz[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                r_t = [[Sequence.GRU._sigmoid(self.br[j] + sum(x_t[b][i] * self.Wr[i][j] for i in range(self.input_dim)) + sum(hs[b][k] * self.Ur[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                ht_tilde = [[math.tanh(self.bh[j] + sum(x_t[b][i] * self.Wh[i][j] for i in range(self.input_dim)) + sum((r_t[b][k] * hs[b][k]) * self.Uh[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                h_t = [[(1.0 - z_t[b][j]) * hs[b][j] + z_t[b][j] * ht_tilde[b][j] for j in range(H)] for b in range(B)]
                hs = h_t
                self._cache["Ts"].append((hs, z_t, r_t, ht_tilde))
                outs.append([[self.by[j] + sum(hs[b][k] * self.Why[k][j] for k in range(H)) for j in range(self.output_dim)] for b in range(B)] if self.output_dim is not None else [row[:] for row in hs])
            if not self.return_sequences:
                last = outs[-1] if outs else ([] if not single else [])
                return last[0] if single else last
            if single: return [o[0] for o in outs]
            return outs

        def backward(self, grad_output):
            Xb = self._cache["X"]
            Ts = self._cache["Ts"]
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            single = not Sequence.GRU._is_3d(Xb)
            z = Sequence.GRU._zeros
            if not self.return_sequences:
                g_last = grad_output if isinstance(grad_output[0], list) else [grad_output]
                G = [z(B, self.output_dim or H) for _ in range(T)]
                G[-1] = g_last
            else:
                G = grad_output
                if single: G = [[g] for g in G]
            # init grads
            D = self.input_dim
            self.grad_Wz = [[0.0] * H for _ in range(D)]
            self.grad_Uz = [[0.0] * H for _ in range(H)]
            self.grad_bz = [0.0] * H
            self.grad_Wr = [[0.0] * H for _ in range(D)]
            self.grad_Ur = [[0.0] * H for _ in range(H)]
            self.grad_br = [0.0] * H
            self.grad_Wh = [[0.0] * H for _ in range(D)]
            self.grad_Uh = [[0.0] * H for _ in range(H)]
            self.grad_bh = [0.0] * H
            if self.output_dim is not None:
                self.grad_Why = [[0.0] * self.output_dim for _ in range(H)]
                self.grad_by = [0.0] * self.output_dim
            dX = [[[0.0] * D for _ in range(T)] for _ in range(B)]
            dh_next = z(B, H)
            for t in reversed(range(T)):
                h_t, z_t, r_t, ht_tilde = Ts[t]
                h_prev = Ts[t - 1][0] if t > 0 else z(B, H)
                # output grad
                if self.output_dim is not None:
                    dy = G[t]
                    dh = z(B, H)
                    for b in range(B):
                        for j in range(self.output_dim):
                            self.grad_by[j] += dy[b][j]
                            for k in range(H): self.grad_Why[k][j] += h_t[b][k] * dy[b][j]
                        for k in range(H):
                            s = 0.0
                            for j in range(self.output_dim): s += dy[b][j] * self.Why[k][j]
                            dh[b][k] = s
                else: dh = G[t]
                for b in range(B):
                    for k in range(H): dh[b][k] += dh_next[b][k]
                # GRU backprop
                dz = z(B, H)
                dht = z(B, H)
                dh_prev = z(B, H)
                for b in range(B):
                    for k in range(H):
                        dz[b][k] = dh[b][k] * (ht_tilde[b][k] - h_prev[b][k])
                        dht[b][k] = dh[b][k] * z_t[b][k]
                        dh_prev[b][k] = dh[b][k] * (1.0 - z_t[b][k])
                # h_tilde backprop
                dh_raw = z(B, H)
                for b in range(B):
                    for k in range(H):
                        dh_raw[b][k] = (1.0 - ht_tilde[b][k] * ht_tilde[b][k]) * dht[b][k]
                        self.grad_bh[k] += dh_raw[b][k]
                # through Uh with r*h_prev
                dr = z(B, H)
                for b in range(B):
                    for j in range(H):
                        # gradient to (r*h_prev) for each j
                        s = 0.0
                        for k in range(H): s += dh_raw[b][k] * self.Uh[j][k]
                        dr[b][j] = s * h_prev[b][j]
                        dh_prev[b][j] += s * r_t[b][j]
                # z and r gate backprop
                dz_raw = Sequence.GRU._zeros(B, H)
                dr_raw = Sequence.GRU._zeros(B, H)
                for b in range(B):
                    for k in range(H):
                        dz_raw[b][k] = dz[b][k] * z_t[b][k] * (1.0 - z_t[b][k])
                        dr_raw[b][k] = dr[b][k] * r_t[b][k] * (1.0 - r_t[b][k])
                        self.grad_bz[k] += dz_raw[b][k]
                        self.grad_br[k] += dr_raw[b][k]
                # grads for weights and inputs
                for b in range(B):
                    x_t = Xb[b][t]
                    for i in range(D):
                        for k in range(H):
                            self.grad_Wz[i][k] += x_t[i] * dz_raw[b][k]
                            self.grad_Wr[i][k] += x_t[i] * dr_raw[b][k]
                            self.grad_Wh[i][k] += x_t[i] * dh_raw[b][k]
                    for j in range(H):
                        for k in range(H):
                            self.grad_Uz[j][k] += h_prev[b][j] * dz_raw[b][k]
                            self.grad_Ur[j][k] += h_prev[b][j] * dr_raw[b][k]
                            self.grad_Uh[j][k] += (r_t[b][j] * h_prev[b][j]) * dh_raw[b][k]
                    # dX
                    for i in range(D):
                        s = 0.0
                        for k in range(H):
                            s += dz_raw[b][k] * self.Wz[i][k]
                            s += dr_raw[b][k] * self.Wr[i][k]
                            s += dh_raw[b][k] * self.Wh[i][k]
                        dX[b][t][i] = s
                    # dh_prev from gate matrices
                    for j in range(H):
                        s = 0.0
                        for k in range(H):
                            s += dz_raw[b][k] * self.Uz[j][k]
                            s += dr_raw[b][k] * self.Ur[j][k]
                        dh_prev[b][j] += s
                dh_next = dh_prev
            if single: return dX[0]
            return dX

        def update(self, lr):
            D, H = self.input_dim, self.hidden_dim
            for i in range(D):
                for j in range(H):
                    self.Wz[i][j] -= lr * self.grad_Wz[i][j]
                    self.Wr[i][j] -= lr * self.grad_Wr[i][j]
                    self.Wh[i][j] -= lr * self.grad_Wh[i][j]
            for i in range(H):
                for j in range(H):
                    self.Uz[i][j] -= lr * self.grad_Uz[i][j]
                    self.Ur[i][j] -= lr * self.grad_Ur[i][j]
                    self.Uh[i][j] -= lr * self.grad_Uh[i][j]
            for j in range(H):
                self.bz[j] -= lr * self.grad_bz[j]
                self.br[j] -= lr * self.grad_br[j]
                self.bh[j] -= lr * self.grad_bh[j]
            if self.output_dim is not None:
                for i in range(H):
                    for j in range(self.output_dim): self.Why[i][j] -= lr * self.grad_Why[i][j]
                for j in range(self.output_dim): self.by[j] -= lr * self.grad_by[j]

    class LSTM:
        """LSTM is Long Short Term Memory.
        
        example:
        >>> from dearning import Sequence
        >>> lstm = Sequence.LSTM(input_dim=3, hidden_dim=5, output_dim=2, return_sequences=False)"""
        def __init__(self, input_dim, hidden_dim, output_dim=None, return_sequences=False):
            self.input_dim = int(input_dim)
            self.hidden_dim = int(hidden_dim)
            self.output_dim = int(output_dim) if output_dim is not None else None
            self.return_sequences = bool(return_sequences)
            H, D = self.hidden_dim, self.input_dim
            # input gate
            self.Wi = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Ui = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bi = [0.0] * H
            # forget gate
            self.Wf = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Uf = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bf = [0.0] * H
            # output gate
            self.Wo = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Uo = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bo = [0.0] * H
            # candidate
            self.Wg = [[random.gauss(0, 1) * math.sqrt(2.0 / D) for _ in range(H)] for _ in range(D)]
            self.Ug = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(H)] for _ in range(H)]
            self.bg = [0.0] * H
            if self.output_dim is not None:
                self.Why = [[random.gauss(0, 1) * math.sqrt(2.0 / H) for _ in range(self.output_dim)] for _ in range(H)]
                self.by = [0.0] * self.output_dim
        def __call__(self, x, training=True): return self.forward(x, training)
        @staticmethod
        def _sigmoid(x): return 1.0 / (1.0 + math.exp(-x))
        @staticmethod
        def _is_3d(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], list)
        @staticmethod
        def _zeros(b, n): return [[0.0] * n for _ in range(b)]
        def _to_batch(self, X):
            if not X: return [], True
            if Sequence.LSTM._is_3d(X): return X, False
            return [X], True

        def forward(self, X, training=True):
            Xb, single = self._to_batch(X)
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            z = Sequence.LSTM._zeros
            h = z(B, H)
            c = z(B, H)
            self._cache = {"X": Xb, "Ts": []}
            outs = []
            for t in range(T):
                x_t = [Xb[b][t] for b in range(B)]
                i_t = [[Sequence.LSTM._sigmoid(self.bi[j] + sum(x_t[b][k] * self.Wi[k][j] for k in range(self.input_dim)) + sum(h[b][k] * self.Ui[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                f_t = [[Sequence.LSTM._sigmoid(self.bf[j] + sum(x_t[b][k] * self.Wf[k][j] for k in range(self.input_dim)) + sum(h[b][k] * self.Uf[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                o_t = [[Sequence.LSTM._sigmoid(self.bo[j] + sum(x_t[b][k] * self.Wo[k][j] for k in range(self.input_dim)) + sum(h[b][k] * self.Uo[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                g_t = [[math.tanh(self.bg[j] + sum(x_t[b][k] * self.Wg[k][j] for k in range(self.input_dim)) + sum(h[b][k] * self.Ug[k][j] for k in range(H))) for j in range(H)] for b in range(B)]
                c = [[f_t[b][j] * c[b][j] + i_t[b][j] * g_t[b][j] for j in range(H)] for b in range(B)]
                h = [[o_t[b][j] * math.tanh(c[b][j]) for j in range(H)] for b in range(B)]
                self._cache["Ts"].append((h, c, i_t, f_t, o_t, g_t))
                outs.append([[self.by[j] + sum(h[b][k] * self.Why[k][j] for k in range(H)) for j in range(self.output_dim)] for b in range(B)] if self.output_dim is not None else [row[:] for row in h])
            if not self.return_sequences:
                last = outs[-1] if outs else ([] if not single else [])
                return last[0] if single else last
            if single: return [o[0] for o in outs]
            return outs

        def backward(self, grad_output):
            Xb = self._cache["X"]
            Ts = self._cache["Ts"]
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            H = self.hidden_dim
            single = not Sequence.LSTM._is_3d(Xb)
            z = Sequence.LSTM._zeros
            if not self.return_sequences:
                g_last = grad_output if isinstance(grad_output[0], list) else [grad_output]
                G = [z(B, self.output_dim or H) for _ in range(T)]
                G[-1] = g_last
            else:
                G = grad_output
                if single: G = [[g] for g in G]
            D = self.input_dim
            self.grad_Wi = [[0.0] * H for _ in range(D)]
            self.grad_Ui = [[0.0] * H for _ in range(H)]
            self.grad_bi = [0.0] * H
            self.grad_Wf = [[0.0] * H for _ in range(D)]
            self.grad_Uf = [[0.0] * H for _ in range(H)]
            self.grad_bf = [0.0] * H
            self.grad_Wo = [[0.0] * H for _ in range(D)]
            self.grad_Uo = [[0.0] * H for _ in range(H)]
            self.grad_bo = [0.0] * H
            self.grad_Wg = [[0.0] * H for _ in range(D)]
            self.grad_Ug = [[0.0] * H for _ in range(H)]
            self.grad_bg = [0.0] * H
            if self.output_dim is not None:
                self.grad_Why = [[0.0] * self.output_dim for _ in range(H)]
                self.grad_by = [0.0] * self.output_dim
            dX = [[[0.0] * D for _ in range(T)] for _ in range(B)]
            dh_next = z(B, H)
            dc_next = z(B, H)
            for t in reversed(range(T)):
                h_t, c_t, i_t, f_t, o_t, g_t = Ts[t]
                h_prev = Ts[t - 1][0] if t > 0 else z(B, H)
                c_prev = Ts[t - 1][1] if t > 0 else z(B, H)
                if self.output_dim is not None:
                    dy = G[t]
                    dh = z(B, H)
                    for b in range(B):
                        for j in range(self.output_dim):
                            self.grad_by[j] += dy[b][j]
                            for k in range(H): self.grad_Why[k][j] += h_t[b][k] * dy[b][j]
                        for k in range(H):
                            s = 0.0
                            for j in range(self.output_dim): s += dy[b][j] * self.Why[k][j]
                            dh[b][k] = s
                else: dh = G[t]
                for b in range(B):
                    for k in range(H): dh[b][k] += dh_next[b][k]
                tanh_c = [[math.tanh(c_t[b][k]) for k in range(H)] for b in range(B)]
                do = z(B, H)
                dc = z(B, H)
                for b in range(B):
                    for k in range(H):
                        do[b][k] = dh[b][k] * tanh_c[b][k]
                        dc[b][k] = dh[b][k] * o_t[b][k] * (1.0 - tanh_c[b][k] * tanh_c[b][k]) + dc_next[b][k]
                di = z(B, H)
                df = z(B, H)
                dg = z(B, H)
                for b in range(B):
                    for k in range(H):
                        di[b][k] = dc[b][k] * g_t[b][k]
                        df[b][k] = dc[b][k] * c_prev[b][k]
                        dg[b][k] = dc[b][k] * i_t[b][k]
                        self.grad_bi[k] += di[b][k] * i_t[b][k] * (1.0 - i_t[b][k])
                        self.grad_bf[k] += df[b][k] * f_t[b][k] * (1.0 - f_t[b][k])
                        self.grad_bo[k] += do[b][k] * o_t[b][k] * (1.0 - o_t[b][k])
                        self.grad_bg[k] += dg[b][k] * (1.0 - g_t[b][k] * g_t[b][k])
                for b in range(B):
                    x_t = Xb[b][t]
                    for i in range(D):
                        for k in range(H):
                            self.grad_Wi[i][k] += x_t[i] * di[b][k] * i_t[b][k] * (1.0 - i_t[b][k])
                            self.grad_Wf[i][k] += x_t[i] * df[b][k] * f_t[b][k] * (1.0 - f_t[b][k])
                            self.grad_Wo[i][k] += x_t[i] * do[b][k] * o_t[b][k] * (1.0 - o_t[b][k])
                            self.grad_Wg[i][k] += x_t[i] * dg[b][k] * (1.0 - g_t[b][k] * g_t[b][k])
                    for j in range(H):
                        for k in range(H):
                            self.grad_Ui[j][k] += h_prev[b][j] * di[b][k] * i_t[b][k] * (1.0 - i_t[b][k])
                            self.grad_Uf[j][k] += h_prev[b][j] * df[b][k] * f_t[b][k] * (1.0 - f_t[b][k])
                            self.grad_Uo[j][k] += h_prev[b][j] * do[b][k] * o_t[b][k] * (1.0 - o_t[b][k])
                            self.grad_Ug[j][k] += h_prev[b][j] * dg[b][k] * (1.0 - g_t[b][k] * g_t[b][k])
                    for i in range(D):
                        s = 0.0
                        for k in range(H):
                            s += (di[b][k] * i_t[b][k] * (1.0 - i_t[b][k])) * self.Wi[i][k]
                            s += (df[b][k] * f_t[b][k] * (1.0 - f_t[b][k])) * self.Wf[i][k]
                            s += (do[b][k] * o_t[b][k] * (1.0 - o_t[b][k])) * self.Wo[i][k]
                            s += (dg[b][k] * (1.0 - g_t[b][k] * g_t[b][k])) * self.Wg[i][k]
                        dX[b][t][i] = s
                dh_prev = z(B, H)
                for b in range(B):
                    for j in range(H):
                        s = 0.0
                        for k in range(H):
                            s += (di[b][k] * i_t[b][k] * (1.0 - i_t[b][k])) * self.Ui[j][k]
                            s += (df[b][k] * f_t[b][k] * (1.0 - f_t[b][k])) * self.Uf[j][k]
                            s += (do[b][k] * o_t[b][k] * (1.0 - o_t[b][k])) * self.Uo[j][k]
                            s += (dg[b][k] * (1.0 - g_t[b][k] * g_t[b][k])) * self.Ug[j][k]
                        dh_prev[b][j] = s
                dh_next = dh_prev
                dc_next = [[dc[b][k] * f_t[b][k] for k in range(H)] for b in range(B)]
            if single: return dX[0]
            return dX

        def update(self, lr):
            D, H = self.input_dim, self.hidden_dim
            for i in range(D):
                for j in range(H):
                    self.Wi[i][j] -= lr * self.grad_Wi[i][j]
                    self.Wf[i][j] -= lr * self.grad_Wf[i][j]
                    self.Wo[i][j] -= lr * self.grad_Wo[i][j]
                    self.Wg[i][j] -= lr * self.grad_Wg[i][j]
            for i in range(H):
                for j in range(H):
                    self.Ui[i][j] -= lr * self.grad_Ui[i][j]
                    self.Uf[i][j] -= lr * self.grad_Uf[i][j]
                    self.Uo[i][j] -= lr * self.grad_Uo[i][j]
                    self.Ug[i][j] -= lr * self.grad_Ug[i][j]
                for j in range(H):
                    self.bi[j] -= lr * self.grad_bi[j]
                    self.bf[j] -= lr * self.grad_bf[j]
                    self.bo[j] -= lr * self.grad_bo[j]
                    self.bg[j] -= lr * self.grad_bg[j]
            if self.output_dim is not None:
                for i in range(H):
                    for j in range(self.output_dim): self.Why[i][j] -= lr * self.grad_Why[i][j]
                for j in range(self.output_dim): self.by[j] -= lr * self.grad_by[j]

    class TCN:
        """TCN is a Temporal Convolutional Network.
        
        example:
        >>> from dearning import Sequence
        >>> tcn = Sequence.TCN(input_dim=3, channels=[16, 32, 64], kernel_size=3, dilations=[1, 2, 4], output_dim=2, return_sequences=False)"""
        def __init__(self, input_dim, channels, kernel_size=3, dilations=None, output_dim=None, return_sequences=False):
            self.input_dim = int(input_dim)
            self.channels = list(channels)
            self.kernel_size = int(kernel_size)
            self.dilations = list(dilations) if dilations is not None else [2 ** i for i in range(len(self.channels))]
            self.output_dim = int(output_dim) if output_dim is not None else None
            self.return_sequences = bool(return_sequences)
            self.blocks = []
            in_ch = self.input_dim
            for out_ch, d in zip(self.channels, self.dilations):
                W = [[[random.gauss(0, 0.1) for _ in range(out_ch)] for _ in range(self.kernel_size)] for _ in range(in_ch)]
                b = [0.0] * out_ch
                W_res = None
                if in_ch != out_ch: W_res = [[random.gauss(0, 0.1) for _ in range(out_ch)] for _ in range(in_ch)]
                self.blocks.append({"W": W, "b": b, "W_res": W_res, "d": int(d)})
                in_ch = out_ch
            if self.output_dim is not None:
                self.W_out = [[random.gauss(0, 1) * math.sqrt(2.0 / in_ch) for _ in range(self.output_dim)] for _ in range(in_ch)]
                self.b_out = [0.0] * self.output_dim
        def __call__(self, x, training=True): return self.forward(x, training)
        @staticmethod
        def _is_3d(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], list)
        @staticmethod
        def _zeros(b, n): return [[0.0] * n for _ in range(b)]
        def _to_batch(self, X):
            if not X: return [], True
            if Sequence.TCN._is_3d(X): return X, False
            return [X], True
        def _conv_causal(self, X, W, b, d):
            T = len(X); C = len(X[0]) if X else 0
            K = self.kernel_size; Cout = len(b)
            Y = [[0.0] * Cout for _ in range(T)]
            for t in range(T):
                for j in range(Cout):
                    s = b[j]
                    for c in range(C):
                        for k in range(K):
                            idx = t - k * d
                            if idx >= 0: s += X[idx][c] * W[c][k][j]
                    Y[t][j] = s
            return Y
        def _conv_causal_backward(self, X, W, dY, d):
            T = len(X); C = len(X[0]) if X else 0
            K = self.kernel_size; Cout = len(dY[0]) if dY else 0
            dX = [[0.0] * C for _ in range(T)]
            dW = [[[0.0] * Cout for _ in range(K)] for _ in range(C)]
            db = [0.0] * Cout
            for t in range(T):
                for j in range(Cout):
                    db[j] += dY[t][j]
                    for c in range(C):
                        for k in range(K):
                            idx = t - k * d
                            if idx >= 0:
                                dW[c][k][j] += X[idx][c] * dY[t][j]
                                dX[idx][c] += W[c][k][j] * dY[t][j]
            return dX, dW, db
        def _residual(self, X, W_res):
            if W_res is None: return X
            T = len(X); Cin = len(X[0]) if X else 0; Cout = len(W_res[0])
            Y = [[0.0] * Cout for _ in range(T)]
            for t in range(T):
                for j in range(Cout):
                    s = 0.0
                    for c in range(Cin): s += X[t][c] * W_res[c][j]
                    Y[t][j] = s
            return Y
        def _residual_backward(self, X, W_res, dY):
            if W_res is None: return dY, None
            T = len(X); Cin = len(X[0]) if X else 0; Cout = len(dY[0]) if dY else 0
            dX = [[0.0] * Cin for _ in range(T)]
            dW = [[0.0] * Cout for _ in range(Cin)]
            for t in range(T):
                for j in range(Cout):
                    for c in range(Cin):
                        dW[c][j] += X[t][c] * dY[t][j]
                        dX[t][c] += W_res[c][j] * dY[t][j]
            return dX, dW

        def forward(self, X, training=True):
            Xb, single = self._to_batch(X)
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            self._cache = {"X": Xb, "blocks": []}
            h = Xb
            for block in self.blocks:
                W, b, W_res, d = block["W"], block["b"], block["W_res"], block["d"]
                pairs = [(h_seq, y_seq, [[a + r for a, r in zip(a_row, r_row)] for a_row, r_row in zip(a_seq, r_seq)]) for h_seq in h for y_seq in 
                 [self._conv_causal(h_seq, W, b, d)] for a_seq in [[[v if v > 0.0 else 0.0 for v in row] for row in y_seq]] for r_seq in [self._residual(h_seq, W_res)]]
                out_block = [p[2] for p in pairs]
                self._cache["blocks"].append([(p[0], p[1]) for p in pairs])
                h = out_block
            self._cache["H_out"] = h
            if self.output_dim is not None: h = [[[self.b_out[j] + sum(h[bi][t][k] * self.W_out[k][j]
             for k in range(len(h[bi][t]))) for j in range(self.output_dim)] for t in range(T)] for bi in range(B)]
            if not self.return_sequences:
                last = h[-1] if not single and h else (h[0] if h else [])
                return last[-1] if not single else last[-1]
            return h[0] if single else h

        def backward(self, grad_output):
            Xb = self._cache["X"]
            B = len(Xb)
            T = len(Xb[0]) if B else 0
            single = not Sequence.TCN._is_3d(Xb)
            if not self.return_sequences:
                g_last = grad_output
                C = len(g_last) if single else len(g_last[0])
                G = [[[0.0] * C for _ in range(T)] for _ in range(B)]
                if single: G[0][-1] = g_last
                else:
                    for b in range(B): G[b][-1] = g_last[b]
            else: G = [grad_output] if single else grad_output
            if self.output_dim is not None:
                self.grad_W_out = [[0.0] * self.output_dim for _ in range(len(self.W_out))]
                self.grad_b_out = [0.0] * self.output_dim
                H_out = self._cache["H_out"]
                dH = [[[0.0] * len(self.W_out) for _ in range(T)] for _ in range(B)]
                for b in range(B):
                    for t in range(T):
                        for j in range(self.output_dim):
                            self.grad_b_out[j] += G[b][t][j]
                            for k in range(len(self.W_out)):
                                self.grad_W_out[k][j] += H_out[b][t][k] * G[b][t][j]
                                dH[b][t][k] += self.W_out[k][j] * G[b][t][j]
                d = dH
            else: d = G
            self.grad_blocks = []
            for block in self.blocks:
                W = block["W"]; b = block["b"]; W_res = block["W_res"]
                C = len(W); K = self.kernel_size; Cout = len(b)
                gW = [[[0.0] * Cout for _ in range(K)] for _ in range(C)]
                gb = [0.0] * Cout
                gW_res = None
                if W_res is not None: gW_res = [[0.0] * len(W_res[0]) for _ in range(len(W_res))]
                self.grad_blocks.append({"W": gW, "b": gb, "W_res": gW_res})
            for bi in reversed(range(len(self.blocks))):
                block = self.blocks[bi]
                gblock = self.grad_blocks[bi]
                W, b, W_res, dila = block["W"], block["b"], block["W_res"], block["d"]
                block_cache = self._cache["blocks"][bi]
                d_new = []
                for bidx in range(B):
                    h_seq, y_seq = block_cache[bidx]
                    d_out = d[bidx]
                    d_act = [[d_out[t][j] * (1.0 if y_seq[t][j] > 0.0 else 0.0) for j in range(len(d_out[t]))] for t in range(T)]
                    dX_conv, dW, db = self._conv_causal_backward(h_seq, W, d_act, dila)
                    for c in range(len(dW)):
                        for k in range(len(dW[c])):
                            for j in range(len(dW[c][k])): gblock["W"][c][k][j] += dW[c][k][j]
                    for j in range(len(db)): gblock["b"][j] += db[j]
                    dX_res, dW_res = self._residual_backward(h_seq, W_res, d_out)
                    if dW_res is not None:
                        for c in range(len(dW_res)):
                            for j in range(len(dW_res[c])): gblock["W_res"][c][j] += dW_res[c][j]
                    d_new.append([[dX_conv[t][c] + dX_res[t][c] for c in range(len(dX_conv[t]))] for t in range(T)])
                d = d_new
            return d[0] if single else d

        def update(self, lr):
            for bi, block in enumerate(self.blocks):
                g = self.grad_blocks[bi]
                W, b, W_res = block["W"], block["b"], block["W_res"]
                for c in range(len(W)):
                    for k in range(len(W[c])):
                        for j in range(len(W[c][k])): W[c][k][j] -= lr * g["W"][c][k][j]
                for j in range(len(b)): b[j] -= lr * g["b"][j]
                if W_res is not None:
                    for c in range(len(W_res)):
                        for j in range(len(W_res[c])): W_res[c][j] -= lr * g["W_res"][c][j]
            if self.output_dim is not None:
                for i in range(len(self.W_out)):
                    for j in range(self.output_dim): self.W_out[i][j] -= lr * self.grad_W_out[i][j]
                for j in range(self.output_dim): self.b_out[j] -= lr * self.grad_b_out[j]

    class Transformer:
        """Transformer.
         
        example:
        >>> from dearning import Sequence
        >>> transformer = Sequence.Transformer(d_model=32, num_heads=4, num_layers=2, d_ff=64, src_vocab_size=1000, tgt_vocab_size=1000, max_len=256, return_sequences=True)"""
        def __init__(self, d_model=32, num_heads=4, num_layers=2, d_ff=64, src_vocab_size=None, tgt_vocab_size=None, max_len=256, return_sequences=True):
            self.d_model = int(d_model)
            self.num_heads = int(num_heads)
            self.num_layers = int(num_layers)
            self.d_ff = int(d_ff)
            self.src_vocab_size = int(src_vocab_size) if src_vocab_size is not None else None
            self.tgt_vocab_size = int(tgt_vocab_size) if tgt_vocab_size is not None else None
            self.max_len = int(max_len)
            self.return_sequences = bool(return_sequences)
            if self.d_model % self.num_heads != 0: raise ValueError("d_model must be divisible by num_heads")
            # embeddings
            if self.src_vocab_size is not None: self.src_embed = [[random.gauss(0, 0.1) for _ in range(self.d_model)] for _ in range(self.src_vocab_size)]
            if self.tgt_vocab_size is not None: self.tgt_embed = [[random.gauss(0, 0.1) for _ in range(self.d_model)] for _ in range(self.tgt_vocab_size)]
            # positional encoding
            self.pos_enc = self._build_positional(self.max_len, self.d_model)
            # encoder/decoder layers
            self.enc_layers = [self._init_layer() for _ in range(self.num_layers)]
            self.dec_layers = [self._init_layer(cross=True) for _ in range(self.num_layers)]
            # output projection
            if self.tgt_vocab_size is not None:
                self.W_out = [[random.gauss(0, 1) * math.sqrt(2.0 / self.d_model) for _ in range(self.tgt_vocab_size)] for _ in range(self.d_model)]
                self.b_out = [0.0] * self.tgt_vocab_size
        def __call__(self, src, tgt, training=True): return self.forward(src, tgt, training)
        @staticmethod
        def _is_token_seq(x): return isinstance(x, list) and x and isinstance(x[0], int)
        @staticmethod
        def _is_token_batch(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], int)
        @staticmethod
        def _is_3d(x): return isinstance(x, list) and x and isinstance(x[0], list) and x[0] and isinstance(x[0][0], list)
        @staticmethod
        def _zeros(b, t, d): return [[[0.0] * d for _ in range(t)] for _ in range(b)]
        def _build_positional(self, max_len, d_model):
            pe = [[0.0] * d_model for _ in range(max_len)]
            for pos in range(max_len):
                for i in range(0, d_model, 2):
                    div = math.exp(-math.log(10000.0) * (i / float(d_model)))
                    pe[pos][i] = math.sin(pos * div)
                    if i + 1 < d_model: pe[pos][i + 1] = math.cos(pos * div)
            return pe
        def _add_positional(self, X):
            B = len(X); T = len(X[0]) if B else 0
            for b in range(B):
                for t in range(T):
                    for i in range(self.d_model): X[b][t][i] += self.pos_enc[t % self.max_len][i]
            return X
        def _embed(self, X, E):
            B = len(X); T = len(X[0]) if B else 0
            out = self._zeros(B, T, self.d_model)
            for b in range(B):
                for t in range(T):
                    idx = int(X[b][t])
                    out[b][t] = E[idx][:]
            return out
        def _to_batch(self, X):
            if not X: return [], True
            if self._is_3d(X): return X, False
            return [X], True
        def _linear_forward(self, X, W, b):
            B = len(X); T = len(X[0]) if B else 0
            Din = len(X[0][0]) if B else 0; Dout = len(b)
            Y = self._zeros(B, T, Dout)
            for bi in range(B):
                for t in range(T):
                    for j in range(Dout):
                        s = b[j]
                        for i in range(Din): s += X[bi][t][i] * W[i][j]
                        Y[bi][t][j] = s
            return Y, (X, W)
        def _linear_backward(self, dY, cache):
            X, W = cache
            B = len(X); T = len(X[0]) if B else 0
            Din = len(X[0][0]) if B else 0; Dout = len(W[0]) if W else 0
            dX = self._zeros(B, T, Din)
            dW = [[0.0] * Dout for _ in range(Din)]
            db = [0.0] * Dout
            for bi in range(B):
                for t in range(T):
                    for j in range(Dout):
                        db[j] += dY[bi][t][j]
                        for i in range(Din):
                            dW[i][j] += X[bi][t][i] * dY[bi][t][j]
                            dX[bi][t][i] += W[i][j] * dY[bi][t][j]
            return dX, dW, db
        def _split_heads(self, X):
            B = len(X); T = len(X[0]) if B else 0
            D = len(X[0][0]) if B else 0; H = self.num_heads
            Dh = D // H
            out = [[[[0.0] * Dh for _ in range(T)] for _ in range(H)] for _ in range(B)]
            for b in range(B):
                for t in range(T):
                    for h in range(H):
                        base = h * Dh
                        out[b][h][t] = X[b][t][base:base + Dh]
            return out

        def _merge_heads(self, X):
            B = len(X); H = len(X[0]) if B else 0
            T = len(X[0][0]) if H else 0; Dh = len(X[0][0][0]) if T else 0
            D = H * Dh
            out = self._zeros(B, T, D)
            for b in range(B):
                for t in range(T):
                    row = []
                    for h in range(H): row += X[b][h][t]
                    out[b][t] = row
            return out
        def _softmax(self, vec):
            if not vec: return vec
            m = max(vec)
            exps = [math.exp(v - m) for v in vec]
            s = sum(exps) or 1.0
            return [e / s for e in exps]
        def _attention(self, Q, K, V, mask=None):
            # Q,K,V: [B][H][T][Dh]
            B = len(Q); H = len(Q[0]) if B else 0
            Tq = len(Q[0][0]) if H else 0; Tk = len(K[0][0]) if H else 0
            Dh = len(Q[0][0][0]) if Tq else 0
            scale = 1.0 / math.sqrt(Dh) if Dh > 0 else 1.0
            scores = [[[[0.0] * Tk for _ in range(Tq)] for _ in range(H)] for _ in range(B)]
            attn = [[[[0.0] * Tk for _ in range(Tq)] for _ in range(H)] for _ in range(B)]
            out = [[[[0.0] * Dh for _ in range(Tq)] for _ in range(H)] for _ in range(B)]
            for b in range(B):
                for h in range(H):
                    for tq in range(Tq):
                        row = [(sum(Q[b][h][tq][d] * K[b][h][tk][d] for d in range(Dh)) * scale) if (mask is None or mask[tq][tk] != 0) else -1e9 for tk in range(Tk)]
                        scores[b][h][tq] = row
                        att = self._softmax(row)
                        attn[b][h][tq] = att
                        out[b][h][tq] = [sum(att[tk] * V[b][h][tk][d] for tk in range(Tk)) for d in range(Dh)]
            return out, (Q, K, V, attn, mask, scale)
        def _attention_backward(self, dOut, cache):
            Q, K, V, attn, mask, scale = cache
            B = len(Q); H = len(Q[0]) if B else 0
            Tq = len(Q[0][0]) if H else 0; Tk = len(K[0][0]) if H else 0
            Dh = len(Q[0][0][0]) if Tq else 0
            dQ = [[[[0.0] * Dh for _ in range(Tq)] for _ in range(H)] for _ in range(B)]
            dK = [[[[0.0] * Dh for _ in range(Tk)] for _ in range(H)] for _ in range(B)]
            dV = [[[[0.0] * Dh for _ in range(Tk)] for _ in range(H)] for _ in range(B)]
            for b in range(B):
                for h in range(H):
                    # dAttn and dV
                    dAttn = [[0.0] * Tk for _ in range(Tq)]
                    for tq in range(Tq):
                        for tk in range(Tk):
                            for d in range(Dh):
                                dAttn[tq][tk] += dOut[b][h][tq][d] * V[b][h][tk][d]
                                dV[b][h][tk][d] += attn[b][h][tq][tk] * dOut[b][h][tq][d]
                    # softmax backward per row
                    dScores = [[0.0] * Tk for _ in range(Tq)]
                    for tq in range(Tq):
                        row = attn[b][h][tq]
                        dot = sum(dAttn[tq][tk] * row[tk] for tk in range(Tk))
                        for tk in range(Tk):
                            dScores[tq][tk] = row[tk] * (dAttn[tq][tk] - dot)
                            if mask is not None and mask[tq][tk] == 0: dScores[tq][tk] = 0.0
                            dScores[tq][tk] *= scale
                    # dQ and dK
                    for tq in range(Tq):
                        for d in range(Dh):
                            dQ[b][h][tq][d] = sum(dScores[tq][tk] * K[b][h][tk][d] for tk in range(Tk))
                    for tk in range(Tk):
                        for d in range(Dh):
                            dK[b][h][tk][d] = sum(dScores[tq][tk] * Q[b][h][tq][d] for tq in range(Tq))
            return dQ, dK, dV
        def _layernorm_forward(self, X, gamma, beta, eps=1e-5):
            B = len(X); T = len(X[0]) if B else 0; D = len(X[0][0]) if B else 0
            Y = self._zeros(B, T, D)
            cache = []
            for b in range(B):
                for t in range(T):
                    row = X[b][t]
                    mean = sum(row) / D
                    var = sum((v - mean) ** 2 for v in row) / D
                    inv = 1.0 / math.sqrt(var + eps)
                    xhat = [(v - mean) * inv for v in row]
                    y = [xhat[i] * gamma[i] + beta[i] for i in range(D)]
                    Y[b][t] = y
                    cache.append((xhat, inv, gamma))
            return Y, (cache, B, T, D, eps)
        def _layernorm_backward(self, dY, cache):
            cache_list, B, T, D, _ = cache
            dX = self._zeros(B, T, D)
            dgamma = [0.0] * D
            dbeta = [0.0] * D
            idx = 0
            for b in range(B):
                for t in range(T):
                    xhat, inv, gamma = cache_list[idx]; idx += 1
                    dy = dY[b][t]
                    for i in range(D):
                        dgamma[i] += dy[i] * xhat[i]
                        dbeta[i] += dy[i]
                    dxhat = [dy[i] * gamma[i] for i in range(D)]
                    sum_dx = sum(dxhat)
                    sum_dx_xhat = sum(dxhat[i] * xhat[i] for i in range(D))
                    for i in range(D):
                        dX[b][t][i] = (dxhat[i] * D - sum_dx - xhat[i] * sum_dx_xhat) * inv / D
            return dX, dgamma, dbeta
        def _ff_forward(self, X, W1, b1, W2, b2):
            h1, c1 = self._linear_forward(X, W1, b1)
            a1 = [[[v if v > 0.0 else 0.0 for v in row] for row in seq] for seq in h1]
            y, c2 = self._linear_forward(a1, W2, b2)
            return y, (c1, a1, c2)
        def _ff_backward(self, dY, cache):
            c1, a1, c2 = cache
            dA1, dW2, db2 = self._linear_backward(dY, c2)
            dH1 = [[[dA1[b][t][i] * (1.0 if a1[b][t][i] > 0.0 else 0.0) for i in range(len(a1[b][t]))] for t in range(len(a1[b]))] for b in range(len(a1))]
            dX, dW1, db1 = self._linear_backward(dH1, c1)
            return dX, dW1, db1, dW2, db2
        def _init_layer(self, cross=False):
            D = self.d_model
            H = self.num_heads
            def init_mha(): return {"Wq": [[random.gauss(0, 0.1) for _ in range(D)] for _ in range(D)], "Wk": [[random.gauss(0, 0.1) for _ in range(D)] for _ in range(D)],
                                    "Wv": [[random.gauss(0, 0.1) for _ in range(D)] for _ in range(D)], "bq": [0.0] * D, "bk": [0.0] * D, "bv": [0.0] * D,
                                    "Wo": [[random.gauss(0, 0.1) for _ in range(D)] for _ in range(D)], "bo": [0.0] * D}
            layer = {"self": init_mha(), "ln1": {"g": [1.0] * D, "b": [0.0] * D},
                "ff": {"W1": [[random.gauss(0, 0.1) for _ in range(self.d_ff)] for _ in range(D)], "b1": [0.0] * self.d_ff,
                       "W2": [[random.gauss(0, 0.1) for _ in range(D)] for _ in range(self.d_ff)], "b2": [0.0] * D},
                "ln2": {"g": [1.0] * D, "b": [0.0] * D}}
            if cross:
                layer["cross"] = init_mha()
                layer["ln3"] = {"g": [1.0] * D, "b": [0.0] * D}
            return layer
        def _mha_forward(self, Xq, Xk, Xv, params, mask=None):
            Q, cQ = self._linear_forward(Xq, params["Wq"], params["bq"])
            K, cK = self._linear_forward(Xk, params["Wk"], params["bk"])
            V, cV = self._linear_forward(Xv, params["Wv"], params["bv"])
            Qh = self._split_heads(Q)
            Kh = self._split_heads(K)
            Vh = self._split_heads(V)
            Oh, attn_cache = self._attention(Qh, Kh, Vh, mask)
            O = self._merge_heads(Oh)
            Y, cO = self._linear_forward(O, params["Wo"], params["bo"])
            cache = (cQ, cK, cV, Qh, Kh, Vh, attn_cache, cO, O)
            return Y, cache
        def _mha_backward(self, dY, params, cache):
            cQ, cK, cV, Qh, Kh, Vh, attn_cache, cO, O = cache
            dO, dWo, dbo = self._linear_backward(dY, cO)
            dOh = self._split_heads(dO)
            dQh, dKh, dVh = self._attention_backward(dOh, attn_cache)
            dQ = self._merge_heads(dQh)
            dK = self._merge_heads(dKh)
            dV = self._merge_heads(dVh)
            dXq, dWq, dbq = self._linear_backward(dQ, cQ)
            dXk, dWk, dbk = self._linear_backward(dK, cK)
            dXv, dWv, dbv = self._linear_backward(dV, cV)
            grads = {"Wq": dWq, "Wk": dWk, "Wv": dWv, "Wo": dWo, "bq": dbq, "bk": dbk, "bv": dbv, "bo": dbo}
            return dXq, dXk, dXv, grads
        def _causal_mask(self, T): return [[1 if j <= i else 0 for j in range(T)] for i in range(T)]

        def forward(self, src, tgt, training=True):
            # prepare inputs
            src_b, src_single = self._to_batch(src); tgt_b, tgt_single = self._to_batch(tgt)
            if self._is_token_batch(src): src_b = self._embed(src_b, self.src_embed)
            elif self._is_token_seq(src): src_b = self._embed([src], self.src_embed)
            if self._is_token_batch(tgt): tgt_b = self._embed(tgt_b, self.tgt_embed)
            elif self._is_token_seq(tgt): tgt_b = self._embed([tgt], self.tgt_embed)
            self._cache = {"src": src_b, "tgt": tgt_b, "enc": [], "dec": []}
            src_b = self._add_positional(src_b); tgt_b = self._add_positional(tgt_b)
            # encoder
            h = src_b
            for layer in self.enc_layers:
                attn, c_attn = self._mha_forward(h, h, h, layer["self"])
                x1 = [[[h[b][t][i] + attn[b][t][i] for i in range(self.d_model)] for t in range(len(h[0]))] for b in range(len(h))]
                ln1, c_ln1 = self._layernorm_forward(x1, layer["ln1"]["g"], layer["ln1"]["b"])
                ff, c_ff = self._ff_forward(ln1, layer["ff"]["W1"], layer["ff"]["b1"], layer["ff"]["W2"], layer["ff"]["b2"])
                x2 = [[[ln1[b][t][i] + ff[b][t][i] for i in range(self.d_model)] for t in range(len(ln1[0]))] for b in range(len(ln1))]
                ln2, c_ln2 = self._layernorm_forward(x2, layer["ln2"]["g"], layer["ln2"]["b"])
                self._cache["enc"].append((h, c_attn, x1, c_ln1, ln1, c_ff, x2, c_ln2)); h = ln2
            enc_out = h
            # decoder
            h = tgt_b; mask = self._causal_mask(len(h[0]) if h else 0)
            for layer in self.dec_layers:
                self_attn, c_sattn = self._mha_forward(h, h, h, layer["self"], mask)
                x1 = [[[h[b][t][i] + self_attn[b][t][i] for i in range(self.d_model)] for t in range(len(h[0]))] for b in range(len(h))]
                ln1, c_ln1 = self._layernorm_forward(x1, layer["ln1"]["g"], layer["ln1"]["b"])
                cross, c_cattn = self._mha_forward(ln1, enc_out, enc_out, layer["cross"])
                x2 = [[[ln1[b][t][i] + cross[b][t][i] for i in range(self.d_model)] for t in range(len(ln1[0]))] for b in range(len(ln1))]
                ln2, c_ln2 = self._layernorm_forward(x2, layer["ln2"]["g"], layer["ln2"]["b"])
                ff, c_ff = self._ff_forward(ln2, layer["ff"]["W1"], layer["ff"]["b1"], layer["ff"]["W2"], layer["ff"]["b2"])
                x3 = [[[ln2[b][t][i] + ff[b][t][i] for i in range(self.d_model)] for t in range(len(ln2[0]))] for b in range(len(ln2))]
                ln3, c_ln3 = self._layernorm_forward(x3, layer["ln3"]["g"], layer["ln3"]["b"])
                self._cache["dec"].append((h, c_sattn, x1, c_ln1, ln1, c_cattn, x2, c_ln2, ln2, c_ff, x3, c_ln3)); h = ln3
            dec_out = h
            # output projection
            if self.tgt_vocab_size is not None: y, c_out = self._linear_forward(dec_out, self.W_out, self.b_out); self._cache["out"] = c_out
            else: y = dec_out; self._cache["out"] = None
            if not self.return_sequences: return y[0][-1] if tgt_single else [seq[-1] for seq in y]
            return y[0] if tgt_single else y

        def backward(self, grad_output):
            src_b = self._cache["src"]
            tgt_b = self._cache["tgt"]
            B = len(tgt_b)
            T = len(tgt_b[0]) if B else 0
            single = not self._is_3d(tgt_b)
            if not self.return_sequences:
                C = len(grad_output) if single else len(grad_output[0])
                G = self._zeros(B, T, C)
                if single: G[0][-1] = grad_output
                else:
                    for b in range(B): G[b][-1] = grad_output[b]
            else: G = [grad_output] if single else grad_output
            # output projection
            if self._cache["out"] is not None: dY, self.grad_W_out, self.grad_b_out = self._linear_backward(G, self._cache["out"])
            else: dY = G
            # decoder backward
            d_dec = dY
            self.grad_dec = []
            d_enc_from_cross = self._zeros(len(src_b), len(src_b[0]) if src_b else 0, self.d_model)
            for li in reversed(range(self.num_layers)):
                layer = self.dec_layers[li]
                h, c_sattn, x1, c_ln1, ln1, c_cattn, x2, c_ln2, ln2, c_ff, x3, c_ln3 = self._cache["dec"][li]
                dx3, dgamma3, dbeta3 = self._layernorm_backward(d_dec, c_ln3)
                dln2 = dx3
                dff, dW1, db1, dW2, db2 = self._ff_backward(dln2, c_ff)
                dx2 = [[[dln2[b][t][i] + dff[b][t][i] for i in range(self.d_model)] for t in range(T)] for b in range(B)]
                dln2n, dgamma2, dbeta2 = self._layernorm_backward(dx2, c_ln2)
                dln1 = dln2n
                dln1a, dln1b, dln1c, grads_c = self._mha_backward(dln1, layer["cross"], c_cattn)
                d_enc_from_cross = [[[d_enc_from_cross[b][t][i] + dln1b[b][t][i] for i in range(self.d_model)] for t in range(len(d_enc_from_cross[0]))] for b in range(len(d_enc_from_cross))]
                dx1 = [[[dln1a[b][t][i] + dln1[b][t][i] for i in range(self.d_model)] for t in range(T)] for b in range(B)]
                dln1n, dgamma1, dbeta1 = self._layernorm_backward(dx1, c_ln1)
                dself_q, dself_k, dself_v, grads_s = self._mha_backward(dln1n, layer["self"], c_sattn)
                d_dec = dself_q
                self.grad_dec.append({"self": grads_s, "cross": grads_c, "ff": {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2},
                    "ln1": {"g": dgamma1, "b": dbeta1}, "ln2": {"g": dgamma2, "b": dbeta2}, "ln3": {"g": dgamma3, "b": dbeta3}})
            self.grad_dec.reverse()
            # encoder backward
            d_enc = d_enc_from_cross
            self.grad_enc = []
            for li in reversed(range(self.num_layers)):
                layer = self.enc_layers[li]
                h, c_attn, x1, c_ln1, ln1, c_ff, x2, c_ln2 = self._cache["enc"][li]
                dx2, dgamma2, dbeta2 = self._layernorm_backward(d_enc, c_ln2)
                dln1 = dx2
                dff, dW1, db1, dW2, db2 = self._ff_backward(dln1, c_ff)
                dx1 = [[[dln1[b][t][i] + dff[b][t][i] for i in range(self.d_model)] for t in range(len(dln1[0]))] for b in range(len(dln1))]
                dln1n, dgamma1, dbeta1 = self._layernorm_backward(dx1, c_ln1)
                dself_q, dself_k, dself_v, grads_s = self._mha_backward(dln1n, layer["self"], c_attn)
                d_enc = dself_q
                self.grad_enc.append({"self": grads_s, "ff": {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2},
                                      "ln1": {"g": dgamma1, "b": dbeta1}, "ln2": {"g": dgamma2, "b": dbeta2}})
            self.grad_enc.reverse()
            # embeddings gradients
            if self.src_vocab_size is not None:
                self.grad_src_embed = [[0.0] * self.d_model for _ in range(self.src_vocab_size)]
            if self.tgt_vocab_size is not None:
                self.grad_tgt_embed = [[0.0] * self.d_model for _ in range(self.tgt_vocab_size)]
            return d_enc, d_dec

        def update(self, lr):
            if self.tgt_vocab_size is not None: # output projection
                for i in range(len(self.W_out)):
                    for j in range(len(self.W_out[0])): self.W_out[i][j] -= lr * self.grad_W_out[i][j]
                for j in range(len(self.b_out)): self.b_out[j] -= lr * self.grad_b_out[j]
            for li in range(self.num_layers): # encoder
                layer = self.enc_layers[li]
                g = self.grad_enc[li]
                for name in ("self",):
                    for key in ("Wq","Wk","Wv","Wo"):
                        W = layer[name][key]; dW = g[name][key]
                        for i in range(len(W)):
                            for j in range(len(W[0])): W[i][j] -= lr * dW[i][j]
                    for key in ("bq","bk","bv","bo"):
                        b = layer[name][key]; db = g[name][key]
                        for j in range(len(b)): b[j] -= lr * db[j]
                for i in range(len(layer["ff"]["W1"])):
                    for j in range(len(layer["ff"]["W1"][0])): layer["ff"]["W1"][i][j] -= lr * g["ff"]["W1"][i][j]
                for j in range(len(layer["ff"]["b1"])): layer["ff"]["b1"][j] -= lr * g["ff"]["b1"][j]
                for i in range(len(layer["ff"]["W2"])):
                    for j in range(len(layer["ff"]["W2"][0])): layer["ff"]["W2"][i][j] -= lr * g["ff"]["W2"][i][j]
                for j in range(len(layer["ff"]["b2"])): layer["ff"]["b2"][j] -= lr * g["ff"]["b2"][j]
                for j in range(len(layer["ln1"]["g"])): layer["ln1"]["g"][j] -= lr * g["ln1"]["g"][j]
                for j in range(len(layer["ln1"]["b"])): layer["ln1"]["b"][j] -= lr * g["ln1"]["b"][j]
                for j in range(len(layer["ln2"]["g"])): layer["ln2"]["g"][j] -= lr * g["ln2"]["g"][j]
                for j in range(len(layer["ln2"]["b"])): layer["ln2"]["b"][j] -= lr * g["ln2"]["b"][j]
            for li in range(self.num_layers): # decoder
                layer = self.dec_layers[li]
                g = self.grad_dec[li]
                for name in ("self","cross"):
                    for key in ("Wq","Wk","Wv","Wo"):
                        W = layer[name][key]; dW = g[name][key]
                        for i in range(len(W)):
                            for j in range(len(W[0])): W[i][j] -= lr * dW[i][j]
                    for key in ("bq","bk","bv","bo"):
                        b = layer[name][key]; db = g[name][key]
                        for j in range(len(b)): b[j] -= lr * db[j]
                for i in range(len(layer["ff"]["W1"])):
                    for j in range(len(layer["ff"]["W1"][0])): layer["ff"]["W1"][i][j] -= lr * g["ff"]["W1"][i][j]
                for j in range(len(layer["ff"]["b1"])): layer["ff"]["b1"][j] -= lr * g["ff"]["b1"][j]
                for i in range(len(layer["ff"]["W2"])):
                    for j in range(len(layer["ff"]["W2"][0])): layer["ff"]["W2"][i][j] -= lr * g["ff"]["W2"][i][j]
                for j in range(len(layer["ff"]["b2"])): layer["ff"]["b2"][j] -= lr * g["ff"]["b2"][j]
                for j in range(len(layer["ln1"]["g"])): layer["ln1"]["g"][j] -= lr * g["ln1"]["g"][j]
                for j in range(len(layer["ln1"]["b"])): layer["ln1"]["b"][j] -= lr * g["ln1"]["b"][j]
                for j in range(len(layer["ln2"]["g"])): layer["ln2"]["g"][j] -= lr * g["ln2"]["g"][j]
                for j in range(len(layer["ln2"]["b"])): layer["ln2"]["b"][j] -= lr * g["ln2"]["b"][j]
                for j in range(len(layer["ln3"]["g"])): layer["ln3"]["g"][j] -= lr * g["ln3"]["g"][j]
                for j in range(len(layer["ln3"]["b"])): layer["ln3"]["b"][j] -= lr * g["ln3"]["b"][j]

class Dense:
    """Dense is a function for a fully connected (linear) neural network layer.
    
    example:
    >>> from dearning import Dense
    >>> layer = Dense(3, 2)  # input_dim=3, output_dim=2"""
    def __init__(self, input_dim, output_dim):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.weights = [[random.gauss(0, 1) * math.sqrt(2.0 / input_dim) for _ in range(output_dim)] for _ in range(input_dim)]
        self.bias = [0.0] * output_dim
    def __call__(self, x, training=True): return self.forward(x, training)

    def forward(self, X, training=True):
        if not X or len(X[0]) != self.input_dim: raise ValueError("Expected input_dim={}, got {}".format(self.input_dim, len(X[0]) if X else None))
        self.input = X
        return [[sum(x[i] * self.weights[i][j] for i in range(self.input_dim)) + self.bias[j] for j in range(self.output_dim)] for x in X]

    def backward(self, grad_output):
        if isinstance(grad_output, (int, float)): grad_output = [[grad_output]]
        B = len(grad_output)
        self.grad_w = [[0.0]*self.output_dim for _ in range(self.input_dim)]
        self.grad_b = [0.0]*self.output_dim
        grad_in = [[0.0]*self.input_dim for _ in range(B)]
        for b in range(B):
            for j in range(self.output_dim):
                if len(grad_output[b]) != len(self.grad_b): raise ValueError("Gradient shape mismatch: got {}, expected {}".format(len(grad_output[b]),len(self.grad_b)))
                self.grad_b[j] += grad_output[b][j]
                for i in range(self.input_dim):
                    self.grad_w[i][j] += self.input[b][i] * grad_output[b][j]
                    grad_in[b][i] += self.weights[i][j] * grad_output[b][j]
        return grad_in

    def update(self, lr):
        for i in range(self.input_dim):
            for j in range(self.output_dim): self.weights[i][j] -= lr * self.grad_w[i][j]
        for j in range(self.output_dim): self.bias[j] -= lr * self.grad_b[j]

class Activation:
    """Activation is a function for the activation layer for relu, sigmoid, tanh, leaky_relu, softmax, and linear.
    
    example:
    >>> from dearning import Activation
    >>> layer = Activation("relu")  # kind can be sigmoid or tanh"""
    def __init__(self, kind="relu", alpha=0.01):
        self.kind = kind
        self.alpha = float(alpha)
    def __call__(self, x, training=True): return self.forward(x, training)

    def _to_2d(x):
        if isinstance(x, (int, float)): return [[x]]
        if isinstance(x, list):
            if not x: return []
            if isinstance(x[0], (int, float)): return [x]
            return x
        raise TypeError("Invalid input shape")
    
    def _restore_shape(x, ref):
        if isinstance(ref, (int, float)): return x[0][0]
        if isinstance(ref, list) and ref and isinstance(ref[0], (int, float)): return x[0]
        return x

    def forward(self, x, training=True):
        self._input_shape = x
        X = Activation._to_2d(x)
        if self.kind == "relu": Y = [[max(0.0, v) for v in row] for row in X]
        elif self.kind == "leaky_relu": Y = [[v if v > 0.0 else self.alpha * v for v in row] for row in X]
        elif self.kind == "sigmoid": Y = [[1.0 / (1.0 + math.exp(-v)) for v in row] for row in X]
        elif self.kind == "tanh": Y = [[math.tanh(v) for v in row] for row in X]
        elif self.kind == "softmax": Y = [list(map(lambda e, s=sum(exps): e / s, exps)) for exps in [[math.exp(v - max(row)) for v in row] for row in X]]
        elif self.kind == "linear": Y = [row[:] for row in X]
        else: raise ValueError("Unknown activation")
        self.output = Y
        return Activation._restore_shape(Y, self._input_shape)
    
    def backward(self, grad_output):
        G = Activation._to_2d(grad_output)
        O = Activation._to_2d(self.output)
        if self.kind == "relu": dX = [[g if o > 0 else 0.0 for g, o in zip(gr, orow)] for gr, orow in zip(G, O)]
        elif self.kind == "leaky_relu": dX = [[g if o > 0 else self.alpha * g for g, o in zip(gr, orow)] for gr, orow in zip(G, O)]
        elif self.kind == "sigmoid": dX = [[g * o * (1 - o) for g, o in zip(gr, orow)] for gr, orow in zip(G, O)]
        elif self.kind == "tanh": dX = [[g * (1 - o * o) for g, o in zip(gr, orow)] for gr, orow in zip(G, O)]
        elif self.kind == "softmax": dX = [[s * (g - sum(gg * ss for gg, ss in zip(gr, srow))) for g, s in zip(gr, srow)] for gr, srow in zip(G, O)]
        elif self.kind == "linear": dX = [row[:] for row in G]
        return Activation._restore_shape(dX, grad_output)

class Dropout:
    """Dropout is a function for regularization layers with random neuron masking during training.
    
    example:
    >>> from dearning import Dropout
    >>> layer = Dropout(rate=0.5)  # rate is the dropout rate"""
    def __init__(self, rate=0.5): self.rate = rate
    def __call__(self, x, training=True): return self.forward(x, training)
    
    def forward(self, x, training=True):
        if not training: return x
        self.mask = [[1.0 if random.random() > self.rate else 0.0 for _ in row] for row in x]
        return [[v*m for v,m in zip(row, mrow)] for row, mrow in zip(x, self.mask)]

    def backward(self, grad_output): return [[g*m for g,m in zip(row, mrow)] for row, mrow in zip(grad_output, self.mask)]

class Merge:
    """Merge is a function for merging two inputs by concatenation.

    example:
    >>> from dearning import Merge
    >>> layer = Merge()"""
    def forward(self, a, b): return [ra + rb for ra, rb in zip(a, b)]
    def backward(self, grad):
        if isinstance(grad, (int, float)): grad = [[grad]]
        if grad and isinstance(grad[0], (int, float)): grad = [grad]
        width = len(grad[0])
        if width % 2 != 0: raise ValueError("Merge backward: cannot split odd feature size")
        mid = width // 2
        g1 = [g[:mid] for g in grad]
        g2 = [g[mid:] for g in grad]
        return g1, g2

class CustomAIModel:
    """CustomAIModel is a function for creating AI models.

    With the following functions:
    - add: adds a layer
    - connect: connects layers in a graph manner
    - forward: performs a forward pass through all layers
    - seed: sets the random seed
    - zero_grad: resets gradients for all layers
    - backward: performs backpropagation on all layers
    - step: updates all model parameters using a specified learning rate.
    
    available losses are only `mse`, `cross_entropy`, and `softmax_cross_entropy`."""
    _add_warned = False
    def __init__(self, loss="mse", name=None):
        self.layers = []
        self.nodes = []
        self.mode = "sequential"
        self.loss = loss
        self.name = name

    def add(self, layer):
        """add is a function to add layers to the model.
        
        example:
        >>> from dearning import CustomAIModel, Dense
        >>> model = CustomAIModel(loss="mse")
        >>> model.add(Dense(1, 1))"""
        if not CustomAIModel._add_warned:
            warnings.warn("CustomAIModel.add() will be deprecated. Use CustomAIModel.connect().", category=DeprecationWarning, stacklevel=2)
            CustomAIModel._add_warned = True
        self.layers.append(layer)

    class InputNode:
        def __init__(self, data):
            self.output = data
            self.grad = None
            self.inputs = []
            self.layer = None

        def __hash__(self): return id(self)

    class Node:
        def __init__(self, layer, inputs):
            self.layer = layer
            self.inputs = inputs
            self.output = None
            self.grad = None
        
        def __hash__(self): return id(self)
        
        def forward_node(self, node):
            if not node.inputs:  return node.output # input node
            xs = [self.forward_node(n) for n in node.inputs]
            return node.layer.forward(xs if len(xs) > 1 else xs[0])
        
        @staticmethod
        def topo_sort(output_node):
            visited = set()
            topo = []
            def visit(n):
                if not isinstance(n, CustomAIModel.Node): return
                if n not in visited:
                    visited.add(n)
                    for i in n.inputs: visit(i)
                    topo.append(n)
            visit(output_node)
            return topo

        def backward_graph(output_node, loss_grad):
            def _ensure_2d(grad):
                if isinstance(grad, (int, float)): return [[grad]]
                if isinstance(grad, list):
                    if grad and isinstance(grad[0], (int, float)): return [grad]
                return grad
            def add_grad(a, b): return [[x + y for x, y in zip(ar, br)] for ar, br in zip(a, b)]
            topo = CustomAIModel.Node.topo_sort(output_node)
            for n in topo: n.grad = None
            output_node.grad = loss_grad
            for n in reversed(topo):
                if not n.layer or not hasattr(n.layer, "backward"): continue
                n.grad = _ensure_2d(n.grad)
                grads = n.layer.backward(n.grad)
                if isinstance(grads, tuple):
                    for parent, g in zip(n.inputs, grads): parent.grad = g if parent.grad is None else add_grad(parent.grad, g)
                else:
                    parent = n.inputs[0]
                    parent.grad = grads if parent.grad is None else add_grad(parent.grad, grads)

    def connect(self, layer, *inputs):
        '''connect is a function to add layers to the model.
        
        example:
        >>> from dearning import CustomAIModel, Dense
        >>> model = CustomAIModel(loss="mse")
        >>> inp = [[1.0, 2.0], [2.0, 3.0]]
        >>> out = model.connect(Dense(2, 1), inp)'''
        self.mode = "graph"
        node_inputs = []
        for inp in inputs:
            if not isinstance(inp, (CustomAIModel.Node, CustomAIModel.InputNode)):
                inp_node = CustomAIModel.InputNode(inp)
                self.nodes.append(inp_node)
                node_inputs.append(inp_node)
            else: node_inputs.append(inp)
        node = CustomAIModel.Node(layer, node_inputs)
        self.nodes.append(node)
        return node

    def forward(self, x):
        for node in self.nodes:
            # InputNode: output sudah ada
            if node.layer is None: continue
            inputs = [i.output for i in node.inputs]
            node.output = node.layer.forward(*inputs)
        return self.nodes[-1].output

    def seed(self, n):
        """seed is a function to set the random seed.
        
        example:
        >>> from dearning import CustomAIModel
        >>> model = CustomAIModel(loss="mse")
        >>> model.seed(42)"""
        random.seed(n)

    def _loss(self, y_pred, y_true):
        total = 0.0
        count = 0
        if self.loss == "mse":
            for yp_row, yt_row in zip(y_pred, y_true):
                for yp, yt in zip(yp_row, yt_row):
                    diff = yp - yt
                    total += diff * diff
                    count += 1
            return total / max(1, count)
        elif self.loss == "cross_entropy":
            epsilon = 1e-8
            for yp_row, yt_row in zip(y_pred, y_true):
                for yp, yt in zip(yp_row, yt_row):
                    yp = min(max(yp, epsilon), 1.0 - epsilon)
                    total += -(yt * math.log(yp) + (1 - yt) * math.log(1 - yp))
                    count += 1
            return total / max(1, count)
        elif self.loss == "softmax_cross_entropy":
            total = 0.0
            count = 0
            for logits, yt_row in zip(y_pred, y_true):
                max_logit = max(logits)
                exps = [math.exp(z - max_logit) for z in logits]
                sum_exps = sum(exps)
                for z, yt in zip(logits, yt_row):
                    if yt == 1:
                        total += -(z - max_logit) + math.log(sum_exps)
                        count += 1
            return total / max(1, count)
        return 0.0

    def _loss_grad(self, y_pred, y_true):
        n = sum(len(row) for row in y_true)
        if self.loss == "mse": return [[2 * (yp - yt) / n for yp, yt in zip(yp_row, yt_row)] for yp_row, yt_row in zip(y_pred, y_true)]
        elif self.loss == "cross_entropy": return [[(yp - yt) / n for yp, yt in zip(yp_row, yt_row)] for yp_row, yt_row in zip(y_pred, y_true)]
        return 0.0
    
    def zero_grad(self):
        """zero_grad is a function to reset grad.
        
        example:
        >>> from dearning import CustomAIModel, Dense
        >>> model = CustomAIModel(loss="mse")
        >>> model.add(Dense(1, 1))
        >>> model.zero_grad()"""
        for layer in self.layers:
            if hasattr(layer, "grad_w"):
                for i in range(len(layer.grad_w)):
                    for j in range(len(layer.grad_w[i])): layer.grad_w[i][j] = 0.0
            if hasattr(layer, "grad_b"): layer.grad_b = [0.0]*len(layer.grad_b)

    def backward(self, grad):
        """backward is a function to run backward propagation on all layers.
        
        example:
        >>> from dearning import CustomAIModel, Dense
        >>> model = CustomAIModel(loss="mse")
        >>> model.add(Dense(1, 1))
        >>> # after forward pass and loss computation
        >>> grad = model._loss_grad(y_pred, y_true)
        >>> model.zero_grad()
        >>> model.backward(grad)"""
        for layer in reversed(self.layers):
            if hasattr(layer, "backward"): grad = layer.backward(grad)
        return grad

    def step(self, lr):
        """step is a function to update all model parameters using a certain learning rate.
        
        example:
        >>> from dearning import CustomAIModel, Dense
        >>> model = CustomAIModel(loss="mse")
        >>> model.add(Dense(1, 1))
        >>> # after backward pass
        >>> lr = 0.1
        >>> model.step(lr)"""
        if self.mode == "graph":
            nodes = self.nodes
            for node in nodes:
                if hasattr(node.layer, "update"): node.layer.update(lr)
        else:
            for layer in self.layers:
                if hasattr(layer, "update"): layer.update(lr)

    def train(self, X, y, epochs=100, learning_rate=0.1, batch_size=None, task=None):
        for _ in range(epochs):
            y_pred = self.forward(X)
            loss = self._loss(y_pred, y)
            grad = self._loss_grad(y_pred, y)
            if self.mode == "graph": self.Node.backward_graph(self.nodes[-1], grad)
            else:
                self.zero_grad()
                self.backward(grad)
            self.step(learning_rate)
        return loss
