update

Author: srush

minitorch/scalar.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional, Sequence, Tuple, Type, Union
+
+import numpy as np
+
+from dataclasses import field
+from .autodiff import Context, Variable, backpropagate, central_difference
+from .scalar_functions import (
+    EQ,
+    LT,
+    Add,
+    Exp,
+    Inv,
+    Log,
+    Mul,
+    Neg,
+    ReLU,
+    ScalarFunction,
+    Sigmoid,
+)
+
+ScalarLike = Union[float, int, "Scalar"]
+
+
+@dataclass
+class ScalarHistory:
+    """`ScalarHistory` stores the history of `Function` operations that was
+    used to construct the current Variable.
+
+    Attributes
+    ----------
+        last_fn : The last Function that was called.
+        ctx : The context for that Function.
+        inputs : The inputs that were given when `last_fn.forward` was called.
+
+    """
+
+    last_fn: Optional[Type[ScalarFunction]] = None
+    ctx: Optional[Context] = None
+    inputs: Sequence[Scalar] = ()
+
+
+# ## Task 1.2 and 1.4
+# Scalar Forward and Backward
+
+_var_count = 0
+
+
+@dataclass
+class Scalar:
+    """A reimplementation of scalar values for autodifferentiation
+    tracking. Scalar Variables behave as close as possible to standard
+    Python numbers while also tracking the operations that led to the
+    number's creation. They can only be manipulated by
+    `ScalarFunction`.
+    """
+
+    data: float
+    history: Optional[ScalarHistory] = field(default_factory=ScalarHistory)
+    derivative: Optional[float] = None
+    name: str = field(default="")
+    unique_id: int = field(default=0)
+
+    def __post_init__(self):
+        global _var_count
+        _var_count += 1
+        object.__setattr__(self, "unique_id", _var_count)
+        object.__setattr__(self, "name", str(self.unique_id))
+        object.__setattr__(self, "data", float(self.data))
+
+    def __repr__(self) -> str:
+        return f"Scalar({self.data})"
+
+    def __mul__(self, b: ScalarLike) -> Scalar:
+        return Mul.apply(self, b)
+
+    def __truediv__(self, b: ScalarLike) -> Scalar:
+        return Mul.apply(self, Inv.apply(b))
+
+    def __rtruediv__(self, b: ScalarLike) -> Scalar:
+        return Mul.apply(b, Inv.apply(self))
+
+    def __bool__(self) -> bool:
+        return bool(self.data)
+
+    def __radd__(self, b: ScalarLike) -> Scalar:
+        return self + b
+
+    def __rmul__(self, b: ScalarLike) -> Scalar:
+        return self * b
+
+    # Variable elements for backprop
+
+    def accumulate_derivative(self, x: Any) -> None:
+        """Add `val` to the the derivative accumulated on this variable.
+        Should only be called during autodifferentiation on leaf variables.
+
+        Args:
+        ----
+            x: value to be accumulated
+
+        """
+        assert self.is_leaf(), "Only leaf variables can have derivatives."
+        if self.derivative is None:
+            self.__setattr__("derivative", 0.0)
+        self.__setattr__("derivative", self.derivative + x)
+
+    def is_leaf(self) -> bool:
+        """True if this variable created by the user (no `last_fn`)"""
+        return self.history is not None and self.history.last_fn is None
+
+    def is_constant(self) -> bool:
+        return self.history is None
+
+    @property
+    def parents(self) -> Iterable[Variable]:
+        """Get the variables used to create this one."""
+        assert self.history is not None
+        return self.history.inputs
+
+    def chain_rule(self, d_output: Any) -> Iterable[Tuple[Variable, Any]]:
+        h = self.history
+        assert h is not None
+        assert h.last_fn is not None
+        assert h.ctx is not None
+
+        raise NotImplementedError("Need to include this file from past assignment.")
+
+    def backward(self, d_output: Optional[float] = None) -> None:
+        """Calls autodiff to fill in the derivatives for the history of this object.
+
+        Args:
+        ----
+            d_output (number, opt): starting derivative to backpropagate through the model
+                                   (typically left out, and assumed to be 1.0).
+
+        """
+        if d_output is None:
+            d_output = 1.0
+        backpropagate(self, d_output)
+
+    raise NotImplementedError("Need to include this file from past assignment.")
+
+
+def derivative_check(f: Any, *scalars: Scalar) -> None:
+    """Checks that autodiff works on a python function.
+    Asserts False if derivative is incorrect.
+
+    Parameters
+    ----------
+        f : function from n-scalars to 1-scalar.
+        *scalars  : n input scalar values.
+
+    """
+    out = f(*scalars)
+    out.backward()
+
+    err_msg = """
+Derivative check at arguments f(%s) and received derivative f'=%f for argument %d,
+but was expecting derivative f'=%f from central difference."""
+    for i, x in enumerate(scalars):
+        check = central_difference(f, *scalars, arg=i)
+        print(str([x.data for x in scalars]), x.derivative, i, check)
+        assert x.derivative is not None
+        np.testing.assert_allclose(
+            x.derivative,
+            check.data,
+            1e-2,
+            1e-2,
+            err_msg=err_msg
+            % (str([x.data for x in scalars]), x.derivative, i, check.data),
+        )

minitorch/scalar_functions.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import minitorch
+
+from . import operators
+from .autodiff import Context
+
+if TYPE_CHECKING:
+    from typing import Tuple
+
+    from .scalar import Scalar, ScalarLike
+
+
+def wrap_tuple(x: float | Tuple[float, ...]) -> Tuple[float, ...]:
+    """Turn a possible value into a tuple"""
+    if isinstance(x, tuple):
+        return x
+    return (x,)
+
+
+class ScalarFunction:
+    """A wrapper for a mathematical function that processes and produces
+    Scalar variables.
+
+    This is a static class and is never instantiated. We use `class`
+    here to group together the `forward` and `backward` code.
+    """
+
+    @classmethod
+    def _backward(cls, ctx: Context, d_out: float) -> Tuple[float, ...]:
+        return wrap_tuple(cls.backward(ctx, d_out))  # type: ignore
+
+    @classmethod
+    def _forward(cls, ctx: Context, *inps: float) -> float:
+        return cls.forward(ctx, *inps)  # type: ignore
+
+    @classmethod
+    def apply(cls, *vals: ScalarLike) -> Scalar:
+        raw_vals = []
+        scalars = []
+        for v in vals:
+            if isinstance(v, minitorch.scalar.Scalar):
+                scalars.append(v)
+                raw_vals.append(v.data)
+            else:
+                scalars.append(minitorch.scalar.Scalar(v))
+                raw_vals.append(v)
+
+        # Create the context.
+        ctx = Context(False)
+
+        # Call forward with the variables.
+        c = cls._forward(ctx, *raw_vals)
+        assert isinstance(c, float), "Expected return type float got %s" % (type(c))
+
+        # Create a new variable from the result with a new history.
+        back = minitorch.scalar.ScalarHistory(cls, ctx, scalars)
+        return minitorch.scalar.Scalar(c, back)
+
+
+# Examples
+class Add(ScalarFunction):
+    """Addition function $f(x, y) = x + y$"""
+
+    @staticmethod
+    def forward(ctx: Context, a: float, b: float) -> float:
+        return a + b
+
+    @staticmethod
+    def backward(ctx: Context, d_output: float) -> Tuple[float, ...]:
+        return d_output, d_output
+
+
+class Log(ScalarFunction):
+    """Log function $f(x) = log(x)$"""
+
+    @staticmethod
+    def forward(ctx: Context, a: float) -> float:
+        ctx.save_for_backward(a)
+        return operators.log(a)
+
+    @staticmethod
+    def backward(ctx: Context, d_output: float) -> float:
+        (a,) = ctx.saved_values
+        return operators.log_back(a, d_output)
+
+
+# To implement.
+
+