Models

The basic interface

A basic model has four OPs, as follows.

class SomeModel[X, Yp, Y, P, Ps, G]:
    def predict(x: X, p: Ps) -> Yp:
        ...

    def loss[n](ys_pred: Tensor[Yp][[n]], ys: Tensor[Y][[n]]) -> G:
        ...

    def update(p: P, g: G) -> P:
        ...

    def learn[n](xs: Tensor[X][[n]], ys: Tensor[Y][[n]], p: Ps, revs: int64) -> Ps:
        ...

Users define predict, loss, and update. Then PyPie generates learn.
predict and update operates on individual elements, while loss is for batched elements. That is, predict takes each X in a Tensor[X][[n]]; update works on each P inside the potentially nested structure Ps.

Here's a concrete example. When subclassing a Model, we omit the op decorators on its methods.

from pypie import Model, Tensor
from typing import Tuple

class Line(Model):
    def predict(x: float32, p: Tuple[float32, float32]) -> float32:
        return p[0] * x + p[1]

    def loss[n](ys_pred: Tensor[float32][[n]], ys: Tensor[float32][[n]]) -> float32:
        return ((ys_pred - ys) ** 2).sum(0)

    def update(p: float32, g: float32) -> float32:
        return p - (0.01 * g)

Here, X is float32, P is float32, and Ps is Tuple[float32, float32].

Training

learn iterates for revs rounds. At each revision, it applies predict to xs and the current p, feeds the result (ys_pred) to loss, computes the gradients, and then generates a new p using update.
learn may take two optional arguments, batch_size and debug:
- if the given batch_size is between 1 and n, then learn creates another loop inside each revision that runs the above process in batches;
- if debug is true, learn also returns the history of all inputs passed to update at every revision.

The stateful interface

Some algorithms need additional information to update parameters, such as Adam and RMS. These algorithms need two more OPs: inflate and deflate.

class SomeStatefulModel[X, Yp, Y, P, Ps, Pi, G]:
    def predict(x: X, p: Ps) -> Yp:
        ...

    def loss[n](ys_pred: Tensor[Yp][[n]], ys: Tensor[Y][[n]]) -> G:
        ...

    def update(s: Pi, g: G) -> Pi:
        ...

    def inflate(p: P) -> Pi:
        ...

    def deflate(s: Pi) -> P:
        ...

    def learn[n](xs: Tensor[X][[n]], ys: Tensor[Y][[n]], p: Ps, revs: int64) -> Ps:
        ...

Now update expects Pi, a leaf parameter inflated with some additional information. inflate initializes this info at the beginning of learn, then deflate removes it at the end of learn.
Here's an example that implements Adam.

from pypie import Model, Tensor
from typing import Tuple

@dataclass
class Params[T]:
    w: T
    b: T

class LineAdam[T](Model):
    def predict(x: T, p: Params[T]) -> T:
        return p.w * x + p.b

    def loss[n](ys_pred: Tensor[T][[n]], ys: Tensor[T][[n]]) -> T:
        return ((ys_pred - ys) ** 2).sum(0)

    def inflate(p: T) -> Tuple[T, T, T]:
        return (p, 0, 0)

    def deflate(s: Tuple[T, T, T]) -> T:
        return s[0]

    def update(s: Tuple[T, T, T], g: T) -> Tuple[T, T, T]:
        m = (0.9 * s[1]) + (0.1 * g)
        v = (0.999 * s[2]) + (0.01 * (g ** 2))
        return (s[0] - (0.05 * (m / (sqrt(v) + 1e-8))), m, v)

Here, T is the leaf parameter P and Tuple[T, T, T] is the inflated leaf Pi.