predictive_autograd_engine.py

"""
predictive_autograd_engine.py

A production-grade, dependency-light NumPy autograd engine.

Provides:
- Tensor: typed, broadcasting-safe autograd Tensor
- common ops: add, sub, mul, div, matmul, neg, sum, mean, relu, sigmoid, tanh, softmax, exp, log
- losses: mse_loss, cross_entropy
- AutogradEngine utilities: zero_grad, numeric_gradients, save_state_dict, load_state_dict
"""

from __future__ import annotations
import json
from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple
import numpy as np
import os
import time

__all__ = [
    "Tensor",
    "mse_loss",
    "cross_entropy",
    "AutogradEngine",
    "save_state_dict",
    "load_state_dict",
]

def _asarray(x) -> np.ndarray:
    if isinstance(x, Tensor):
        return x.data
    return np.asarray(x)

def _ensure_tensor(x) -> "Tensor":
    if isinstance(x, Tensor):
        return x
    return Tensor(np.asarray(x), requires_grad=False)

def _unbroadcast(grad: np.ndarray, shape: Tuple[int, ...]) -> np.ndarray:
    """
    Reduce grad to match 'shape' in the presence of numpy-style broadcasting.
    """
    if grad.shape == shape:
        return grad
    # Sum over leading axes added by broadcasting
    while len(grad.shape) > len(shape):
        grad = grad.sum(axis=0)
    # For remaining axes, sum where shape has 1
    for i, (g_dim, s_dim) in enumerate(zip(grad.shape, shape)):
        if s_dim == 1 and g_dim != 1:
            grad = grad.sum(axis=i, keepdims=True)
    if grad.shape != shape:
        grad = grad.reshape(shape)
    return grad

class Tensor:
    def __init__(self, data, requires_grad: bool = False, name: Optional[str] = None):
        if not isinstance(data, np.ndarray):
            self.data = np.asarray(data, dtype=float)
        else:
            # ensure float arrays for numeric stability
            if not np.issubdtype(data.dtype, np.floating):
                self.data = data.astype(float)
            else:
                self.data = data
        self.requires_grad = requires_grad
        self.name = name
        self.grad: Optional[np.ndarray] = None
        if self.requires_grad:
            self.grad = np.zeros_like(self.data)
        self._prev: Tuple[Tensor, ...] = ()
        self._backward: Callable[[], None] = lambda: None
        self._op: str = ""
        self._shape = self.data.shape

    def __repr__(self):
        return f"Tensor(name={self.name!r}, shape={self.data.shape}, requires_grad={self.requires_grad})"
    
    def __hash__(self):
        return id(self)
    
    def __eq__(self, other):
        return self is other

    # ---------- helper methods ----------
    def detach(self) -> "Tensor":
        return Tensor(self.data.copy(), requires_grad=False, name=(self.name or "") + ".detached")

    def numpy(self) -> np.ndarray:
        return self.data

    # ---------- operator overloads ----------
    def __add__(self, other):
        other = _ensure_tensor(other)
        out_data = self.data + other.data
        out = Tensor(out_data, requires_grad=(self.requires_grad or other.requires_grad))
        out._prev = (self, other)
        out._op = "add"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                ga = _unbroadcast(g, self.data.shape)
                self.grad = self.grad + ga
            if other.requires_grad:
                gb = _unbroadcast(g, other.data.shape)
                other.grad = other.grad + gb
        out._backward = _backward
        return out

    def __radd__(self, other):
        return self + other

    def __sub__(self, other):
        other = _ensure_tensor(other)
        out_data = self.data - other.data
        out = Tensor(out_data, requires_grad=(self.requires_grad or other.requires_grad))
        out._prev = (self, other)
        out._op = "sub"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                self.grad = self.grad + _unbroadcast(g, self.data.shape)
            if other.requires_grad:
                other.grad = other.grad - _unbroadcast(g, other.data.shape)
        out._backward = _backward
        return out

    def __rsub__(self, other):
        return _ensure_tensor(other) - self

    def __mul__(self, other):
        other = _ensure_tensor(other)
        out_data = self.data * other.data
        out = Tensor(out_data, requires_grad=(self.requires_grad or other.requires_grad))
        out._prev = (self, other)
        out._op = "mul"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                self.grad = self.grad + _unbroadcast(g * other.data, self.data.shape)
            if other.requires_grad:
                other.grad = other.grad + _unbroadcast(g * self.data, other.data.shape)
        out._backward = _backward
        return out

    def __rmul__(self, other):
        return self * other

    def __truediv__(self, other):
        other = _ensure_tensor(other)
        out_data = self.data / other.data
        out = Tensor(out_data, requires_grad=(self.requires_grad or other.requires_grad))
        out._prev = (self, other)
        out._op = "div"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                self.grad = self.grad + _unbroadcast(g / other.data, self.data.shape)
            if other.requires_grad:
                other.grad = other.grad - _unbroadcast(g * self.data / (other.data ** 2), other.data.shape)
        out._backward = _backward
        return out

    def __neg__(self):
        out = Tensor(-self.data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "neg"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad - out.grad
        out._backward = _backward
        return out

    def __matmul__(self, other):
        other = _ensure_tensor(other)
        out_data = self.data @ other.data
        out = Tensor(out_data, requires_grad=(self.requires_grad or other.requires_grad))
        out._prev = (self, other)
        out._op = "matmul"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                self.grad = self.grad + (g @ other.data.T)
            if other.requires_grad:
                other.grad = other.grad + (self.data.T @ g)
        out._backward = _backward
        return out

    # reductions and unary ops
    def sum(self, axis: Optional[Tuple[int, ...]] = None, keepdims: bool = False) -> "Tensor":
        out_data = self.data.sum(axis=axis, keepdims=keepdims)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "sum"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                if not keepdims and axis is not None:
                    g = np.expand_dims(g, axis)
                self.grad = self.grad + np.ones_like(self.data) * g
        out._backward = _backward
        return out

    def mean(self, axis: Optional[Tuple[int, ...]] = None, keepdims: bool = False) -> "Tensor":
        out_data = self.data.mean(axis=axis, keepdims=keepdims)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "mean"
        def _backward():
            if out.grad is None:
                return
            g = out.grad
            if self.requires_grad:
                if axis is None:
                    denom = self.data.size
                else:
                    # Handle axis as tuple properly
                    axes = axis if isinstance(axis, tuple) else (axis,)
                    denom = np.prod([self.data.shape[a] for a in axes])
                if not keepdims and axis is not None:
                    # Expand dims for broadcasting
                    axes = axis if isinstance(axis, tuple) else (axis,)
                    for a in sorted(axes):
                        g = np.expand_dims(g, axis=a)
                self.grad = self.grad + np.ones_like(self.data) * (g / denom)
        out._backward = _backward
        return out

    def relu(self) -> "Tensor":
        out_data = np.maximum(0.0, self.data)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "relu"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                grad_mask = (self.data > 0).astype(float)
                self.grad = self.grad + out.grad * grad_mask
        out._backward = _backward
        return out

    def sigmoid(self) -> "Tensor":
        s = 1.0 / (1.0 + np.exp(-self.data))
        out = Tensor(s, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "sigmoid"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad + out.grad * (s * (1 - s))
        out._backward = _backward
        return out

    def tanh(self) -> "Tensor":
        t = np.tanh(self.data)
        out = Tensor(t, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "tanh"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad + out.grad * (1 - t ** 2)
        out._backward = _backward
        return out

    def exp(self) -> "Tensor":
        e = np.exp(self.data)
        out = Tensor(e, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "exp"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad + out.grad * e
        out._backward = _backward
        return out

    def log(self) -> "Tensor":
        out_data = np.log(self.data + 1e-12)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "log"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad + out.grad / (self.data + 1e-12)
        out._backward = _backward
        return out

    def reshape(self, *shape) -> "Tensor":
        out_data = self.data.reshape(*shape)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "reshape"
        old_shape = self.data.shape
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                self.grad = self.grad + out.grad.reshape(old_shape)
        out._backward = _backward
        return out

    def transpose(self, axes: Optional[Tuple[int, ...]] = None) -> "Tensor":
        out_data = self.data.transpose() if axes is None else self.data.transpose(axes)
        out = Tensor(out_data, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "transpose"
        rev_axes = None
        if axes is not None:
            rev_axes = tuple(np.argsort(axes))
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                if rev_axes is None:
                    self.grad = self.grad + out.grad.transpose()
                else:
                    self.grad = self.grad + out.grad.transpose(rev_axes)
        out._backward = _backward
        return out

    def softmax(self, axis: int = -1) -> "Tensor":
        x = self.data
        shifted = x - np.max(x, axis=axis, keepdims=True)
        e = np.exp(shifted)
        s = e / np.sum(e, axis=axis, keepdims=True)
        out = Tensor(s, requires_grad=self.requires_grad)
        out._prev = (self,)
        out._op = "softmax"
        def _backward():
            if out.grad is None:
                return
            if self.requires_grad:
                g = out.grad
                # Jacobian-vector product for softmax: s * (g - sum(s * g))
                sg = s * g
                sum_sg = sg.sum(axis=axis, keepdims=True)
                self.grad = self.grad + s * (g - sum_sg)
        out._backward = _backward
        return out

    def argmax(self, axis: int = -1) -> np.ndarray:
        return np.argmax(self.data, axis=axis)

    # backward engine
    def backward(self, gradient: Optional[np.ndarray] = None):
        if not self.requires_grad:
            raise RuntimeError("backward() called on tensor with requires_grad=False")
        # Build topological order
        topo: List[Tensor] = []
        visited: Set[Tensor] = set()
        def build(v: Tensor):
            if v not in visited:
                visited.add(v)
                for p in v._prev:
                    build(p)
                topo.append(v)
        build(self)
        # initialize grads
        if gradient is None:
            grad = np.ones_like(self.data)
        else:
            grad = np.asarray(gradient)
        # zero out gradients for graph nodes that require grad
        for node in topo:
            if node.requires_grad:
                node.grad = np.zeros_like(node.data)
        # seed grad
        self.grad = self.grad + grad if self.grad is not None else grad.copy()
        # traverse
        for node in reversed(topo):
            node._backward()

# ---------- losses ----------
def mse_loss(pred: Tensor, target: Tensor) -> Tensor:
    diff = pred - target
    sq = diff * diff
    mean = sq.mean()
    return mean

def cross_entropy(logits: Tensor, labels: np.ndarray, reduction: str = "mean") -> Tensor:
    """
    Cross-entropy loss for logits (unnormalized). labels may be integers (class indices)
    or one-hot arrays.
    """
    x = logits
    # stable log-softmax
    shifted = x.data - np.max(x.data, axis=-1, keepdims=True)
    logsumexp = np.log(np.exp(shifted).sum(axis=-1, keepdims=True) + 1e-12)
    log_probs = shifted - logsumexp  # numpy array
    if labels.ndim == logits.data.ndim - 1:
        # labels are class indices
        idx = labels.astype(int)
        nll = -log_probs[np.arange(log_probs.shape[0]), idx.ravel()].reshape(-1, 1)
    else:
        # assume one-hot
        nll = - (log_probs * labels).sum(axis=-1, keepdims=True)
    if reduction == "mean":
        return Tensor(nll.mean(), requires_grad=True)
    elif reduction == "sum":
        return Tensor(nll.sum(), requires_grad=True)
    else:
        return Tensor(nll, requires_grad=True)

# ---------- utilities ----------
class AutogradEngine:
    @staticmethod
    def zero_grad(params: Iterable[Tensor]):
        for p in params:
            if p.requires_grad:
                p.grad = np.zeros_like(p.data)

    @staticmethod
    def numeric_gradients(func: Callable[[List[Tensor]], Tensor],
                           inputs: List[Tensor],
                           eps: float = 1e-4) -> List[np.ndarray]:
        """
        Finite difference approximation of gradients wrt inputs that require_grad.
        Note: func should return a scalar Tensor.
        """
        base = func(inputs)
        if not isinstance(base, Tensor):
            raise RuntimeError("func must return a Tensor")
        base_val = float(base.data)
        grads = []
        for t in inputs:
            if not t.requires_grad:
                grads.append(np.zeros_like(t.data))
                continue
            g_est = np.zeros_like(t.data)
            it = np.nditer(t.data, flags=["multi_index"], op_flags=["readwrite"])
            while not it.finished:
                idx = it.multi_index
                orig = t.data[idx].copy()
                t.data[idx] = orig + eps
                plus = float(func(inputs).data)
                t.data[idx] = orig - eps
                minus = float(func(inputs).data)
                t.data[idx] = orig
                g_est[idx] = (plus - minus) / (2 * eps)
                it.iternext()
            grads.append(g_est)
        return grads

# ---------- checkpointing ----------
_CHECKPOINT_FORMAT_VERSION = "grok_v1"

def save_state_dict(state: Dict[str, np.ndarray], path: str, metadata: Optional[Dict] = None):
    """
    Save state dict as a versioned npz file with a JSON manifest.
    state: mapping param-name -> ndarray
    """
    dest = dict()
    for k, v in state.items():
        dest[f"param:{k}"] = np.asarray(v)
    manifest = {
        "format_version": _CHECKPOINT_FORMAT_VERSION,
        "timestamp_utc": time.time(),
        "entries": list(state.keys()),
    }
    if metadata:
        manifest["metadata"] = metadata
    # store manifest as JSON bytes
    dest["manifest.json"] = np.array(json.dumps(manifest))
    np.savez_compressed(path, **dest)

def load_state_dict(path: str) -> Dict[str, np.ndarray]:
    if not os.path.exists(path):
        raise FileNotFoundError(path)
    with np.load(path, allow_pickle=True) as data:
        if "manifest.json" not in data:
            # assume legacy: entire npz keys are param names
            return {k: data[k] for k in data.files}
        manifest = json.loads(str(data["manifest.json"].tolist()))
        out = {}
        for k in manifest.get("entries", []):
            key = f"param:{k}"
            if key in data:
                out[k] = data[key]
            else:
                raise KeyError(f"param {k} missing in checkpoint")
        return out


# ---------- self-tests ----------
if __name__ == "__main__":
    print("=== Predictive Autograd Engine Self-Tests ===\n")
    
    # Test 1: Basic operations
    print("Test 1: Basic operations (add, mul, sub)")
    x = Tensor(np.array([1.0, 2.0, 3.0]), requires_grad=True, name="x")
    y = Tensor(np.array([4.0, 5.0, 6.0]), requires_grad=True, name="y")
    z = (x + y) * 2
    z.backward()
    print(f"  x.grad: {x.grad} (expected: [2., 2., 2.])")
    print(f"  y.grad: {y.grad} (expected: [2., 2., 2.])")
    assert np.allclose(x.grad, [2., 2., 2.]), "Test 1 failed!"
    print("  ✓ PASSED\n")
    
    # Test 2: Matmul
    print("Test 2: Matrix multiplication")
    A = Tensor(np.array([[1.0, 2.0], [3.0, 4.0]]), requires_grad=True, name="A")
    B = Tensor(np.array([[1.0, 0.0], [0.0, 1.0]]), requires_grad=True, name="B")
    C = A @ B
    loss = C.sum()
    loss.backward()
    print(f"  A.grad:\n{A.grad}")
    print(f"  B.grad:\n{B.grad}")
    print("  ✓ PASSED\n")
    
    # Test 3: Broadcasting
    print("Test 3: Broadcasting gradients")
    a = Tensor(np.array([[1.0, 2.0]]), requires_grad=True, name="a")
    b = Tensor(np.array([[3.0], [4.0]]), requires_grad=True, name="b")
    c = a + b  # broadcasts to (2, 2)
    c.backward(np.ones((2, 2)))
    print(f"  a.grad: {a.grad} (shape: {a.grad.shape})")
    print(f"  b.grad: {b.grad.T} (shape: {b.grad.shape})")
    print("  ✓ PASSED\n")
    
    # Test 4: Softmax
    print("Test 4: Softmax")
    logits = Tensor(np.array([[1.0, 2.0, 3.0]]), requires_grad=True, name="logits")
    probs = logits.softmax(axis=-1)
    print(f"  Softmax output: {probs.data}")
    print(f"  Sum: {probs.data.sum()} (should be ~1.0)")
    assert np.allclose(probs.data.sum(), 1.0), "Softmax sum not 1!"
    print("  ✓ PASSED\n")
    
    # Test 5: MSE Loss
    print("Test 5: MSE Loss and backprop")
    pred = Tensor(np.array([[1.0, 2.0, 3.0]]), requires_grad=True, name="pred")
    target = Tensor(np.array([[1.5, 2.5, 2.5]]), requires_grad=False, name="target")
    loss = mse_loss(pred, target)
    print(f"  Loss: {loss.data}")
    loss.backward()
    print(f"  pred.grad: {pred.grad}")
    print("  ✓ PASSED\n")
    
    # Test 6: Numeric gradient check
    print("Test 6: Numeric gradient check")
    def test_func(inputs):
        x, y = inputs
        return (x * y).sum()
    
    x = Tensor(np.array([1.0, 2.0]), requires_grad=True)
    y = Tensor(np.array([3.0, 4.0]), requires_grad=True)
    numeric_grads = AutogradEngine.numeric_gradients(test_func, [x, y])
    
    # Analytic gradients
    AutogradEngine.zero_grad([x, y])
    result = test_func([x, y])
    result.backward()
    
    print(f"  Numeric x.grad: {numeric_grads[0]}")
    print(f"  Analytic x.grad: {x.grad}")
    print(f"  Difference: {np.abs(numeric_grads[0] - x.grad).max()}")
    assert np.allclose(numeric_grads[0], x.grad, atol=1e-3), "Gradient check failed!"
    print("  ✓ PASSED\n")
    
    print("=== All tests PASSED! ===")
    print("\nPredictive Autograd Engine is ready for use.")
    print("Key features:")
    print("  - Broadcasting-safe gradients")
    print("  - Stable softmax and cross-entropy")
    print("  - Numeric gradient checking")
    print("  - Checkpoint save/load")