def get_batch() -> Tuple[Tensor[float32, D32, D4], Tensor[float32, D32, D1]]:
"""Builds a batch i.e. (x, f(x)) pair."""
batch_size = 32
random: Tensor[float32, D32] = torch.randn(batch_size)
x = make_features(random)
y = f(x)
return x, y
from itertools import count
from typing import Sequence, Tuple, TypeVar
import _torch as torch
import _torch.nn.functional as F
from _torch import Tensor, float32
from typing_extensions import Literal
DType = TypeVar("DType", int, float)
POLY_DEGREE: int = 4
D1 = Literal[1]
D4 = Literal[4]
D32 = Literal[32]
W_target: Tensor[float32, [D4, D1]] = torch.randn(POLY_DEGREE, 1) * 5
b_target: Tensor[float32, [D1]] = torch.randn(1) * 5
N = TypeVar("N")
def make_features(x: Tensor[DType, [N]]) -> Tensor[DType, [N, D4]]:
"""Builds features i.e. a matrix with columns [x, x^2, x^3, x^4]."""
# x = x.unsqueeze(1)
x2 = torch.unsqueeze(x, 1)
# return torch.cat([x ** i for i in range(1, POLY_DEGREE+1)], 1)
r: Tensor[DType,
[N, D4]] = torch.cat([x2**i for i in range(1, POLY_DEGREE + 1)],
1)
return r