PyTorch for Temporal AI

"Give me a tensor, a loss, and a gradient that flows backward through time, and I will fit anything you can shape into a batch. Forget the gradient and I will silently fit your future to your past, and call the leak an accuracy."

A Backpropagation Graph That Has Seen Your Validation Set

1. Tensors and Autograd

A PyTorch tensor is a multidimensional array that also remembers, when asked, how it was computed, so the framework can replay that computation backward to produce gradients. Creating tensors is direct: you can build them from Python lists or NumPy arrays, or allocate them with shape constructors. Every tensor carries three properties you will check constantly: its shape (the size along each axis), its dtype (the numeric type, usually float32 for model weights), and its device (CPU or a specific GPU). Getting these three right prevents the large majority of runtime errors in temporal models, where a transposed time axis or a stray float64 is the usual culprit.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])      # from a Python list, dtype inferred as float32
z = torch.zeros(2, 3)                             # shape constructor: a 2x3 tensor of zeros
r = torch.randn(4, 16, 8)                         # random normal, our (batch, time, features) shape

print(x.shape, x.dtype, x.device)                # torch.Size([2, 2]) torch.float32 cpu
print(r.shape)                                   # torch.Size([4, 16, 8]): 4 series, 16 steps, 8 features

# Move to GPU when one is present; keep a single device variable for the whole program
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
r = r.to(device)                                 # tensors must share a device to interact

Autograd is the engine that makes learning possible. When a tensor is created with requires_grad=True, PyTorch records every operation applied to it into a dynamic computation graph. Calling .backward() on a scalar (typically the loss) walks that graph in reverse, accumulating the partial derivative of the scalar with respect to each leaf tensor into that tensor's .grad field. For a scalar loss $L$ and a parameter $\theta$, autograd computes $\partial L / \partial \theta$ exactly, by the chain rule, with no manual derivation. The example below differentiates a simple quadratic so you can check the result by hand: for $y = \sum (w x)^2$ the gradient with respect to $w$ is $\sum 2 x^2 w$.

w = torch.tensor([2.0, 3.0], requires_grad=True)  # a leaf tensor we want gradients for
x = torch.tensor([1.0, 4.0])                       # data, no gradient tracking needed
y = ((w * x) ** 2).sum()                           # scalar output: builds the autograd graph

y.backward()                                       # reverse-mode pass: fills w.grad
print(w.grad)                                      # tensor([4., 96.]) == 2 * x**2 * w, checked by hand

w.grad.zero_()                                     # gradients accumulate, so zero them before the next step

2. nn.Module: Defining Models

An nn.Module is the base class for every model and every layer. You subclass it, register submodules and parameters in __init__, and define the forward computation in forward. Any tensor you wrap in nn.Parameter (and every submodule you assign as an attribute) is tracked automatically, so model.parameters() returns the full set of learnable tensors for the optimizer, and model.to(device) moves them all at once. The minimal module below defines a linear forecaster from scratch so you can see what the framework registers; in practice you would assemble built-in layers rather than declare raw parameters.

import torch.nn as nn

class LinearForecaster(nn.Module):
    def __init__(self, n_features, horizon):
        super().__init__()                              # always call the base constructor first
        self.W = nn.Parameter(torch.randn(n_features, horizon) * 0.01)  # tracked weight
        self.b = nn.Parameter(torch.zeros(horizon))                    # tracked bias

    def forward(self, x):                               # x: (batch, features) of the last step
        return x @ self.W + self.b                      # (batch, horizon) point forecast

model = LinearForecaster(n_features=8, horizon=4)
print(sum(p.numel() for p in model.parameters()))      # 36 = 8*4 weights + 4 biases

You will rarely write recurrences or convolutions by hand, because PyTorch ships the sequence layers that Part III studies in depth. The recurrent modules nn.RNN, nn.LSTM, and nn.GRU (built in detail in Chapter 10) consume a $(B, T, F)$ batch and return the per-step outputs plus the final hidden state. nn.Conv1d is the building block of the temporal convolutional networks of Chapter 11, and it expects channels before time, shape $(B, F, T)$, so a transpose is usually needed. nn.TransformerEncoderLayer and nn.MultiheadAttention back the architectures of Chapter 12. The block below instantiates each so you can see its input and output shapes side by side.

B, T, F = 4, 16, 8
x = torch.randn(B, T, F)                                # (batch, time, features)

lstm = nn.LSTM(input_size=F, hidden_size=32, num_layers=2, batch_first=True)
out, (h_n, c_n) = lstm(x)                               # out: (B, T, 32); h_n: (2, B, 32) per layer
last = out[:, -1, :]                                    # (B, 32): summary at the final step

conv = nn.Conv1d(in_channels=F, out_channels=32, kernel_size=3, padding=1)
y = conv(x.transpose(1, 2))                             # Conv1d wants (B, F, T); transpose then back
y = y.transpose(1, 2)                                   # (B, T, 32)

enc_layer = nn.TransformerEncoderLayer(d_model=F, nhead=2, batch_first=True)
encoder = nn.TransformerEncoder(enc_layer, num_layers=2)
z = encoder(x)                                          # (B, T, F): attention over the time axis

3. Datasets and DataLoaders for Time Series

A Dataset answers two questions: how many samples exist (__len__) and how to fetch the $i$-th one (__getitem__). For time series the unit sample is a sliding window: a contiguous block of past observations as input, paired with the value or values that follow as the target. The single rule that governs this, restated from the leakage discipline of Chapter 2.5, is that the target window must lie strictly after the input window in time, and that any scaling statistics must be fit on training data alone. The windowing dataset below enforces the temporal gap by construction: index $i$ reads inputs from $[i, i+L)$ and the target from $[i+L, i+L+H)$, so the target is always in the future.

from torch.utils.data import Dataset, DataLoader

class WindowDataset(Dataset):
    def __init__(self, series, lookback, horizon):
        self.series = torch.as_tensor(series, dtype=torch.float32)  # (N, F)
        self.L, self.H = lookback, horizon

    def __len__(self):
        return len(self.series) - self.L - self.H + 1   # number of valid windows

    def __getitem__(self, i):
        x = self.series[i : i + self.L]                 # input window: steps [i, i+L)
        y = self.series[i + self.L : i + self.L + self.H, 0]   # target: steps [i+L, i+L+H), first feature
        return x, y                                     # x: (L, F)   y: (H,); target strictly after input

The DataLoader wraps a dataset to produce batches, optionally shuffled, optionally across several worker processes. It stacks the per-sample $(L, F)$ inputs into a $(B, L, F)$ batch and the $(H,)$ targets into $(B, H)$, exactly the convention of Section 1. Shuffling the order of windows is fine and helps optimization; it does not break the temporal order inside each window, which is fixed. When sequences have different lengths (variable-length events, ragged clinical records), you cannot stack them directly, so PyTorch offers two tools: pad_sequence to pad a batch to a common length, and pack_padded_sequence to tell a recurrent module to skip the padding so it does not corrupt the hidden state.

from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence

train = WindowDataset(train_series, lookback=16, horizon=4)
loader = DataLoader(train, batch_size=64, shuffle=True, num_workers=2, drop_last=True)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)                               # torch.Size([64, 16, F]) torch.Size([64, 4])

# Variable-length sequences: pad, then pack so the LSTM ignores the padded steps
seqs = [torch.randn(n, F) for n in (12, 7, 19)]         # three sequences of different lengths
lengths = torch.tensor([s.size(0) for s in seqs])
padded = pad_sequence(seqs, batch_first=True)           # (3, 19, F), shorter ones zero-padded
packed = pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False)

4. The Training and Evaluation Loop

Training is a loop over epochs, each epoch a loop over batches. For every batch you perform five steps in a fixed order: zero the gradients, run the forward pass, compute the loss, call backward, and step the optimizer. Two additions matter for temporal models. Mixed precision (via torch.amp.autocast and a GradScaler) runs the forward pass in half precision for speed and memory, which lets you fit longer sequences. Gradient clipping caps the global gradient norm before the optimizer step, which tames the exploding gradients that recurrent networks are prone to (a problem dissected in Chapter 10). The loop below includes both, plus checkpointing of the best model by validation loss.

import torch.nn as nn

model = model.to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
loss_fn = nn.MSELoss()
scaler = torch.amp.GradScaler(device.type)             # scales the loss for stable fp16 gradients
best_val = float("inf")

for epoch in range(50):
    model.train()                                      # enable dropout and batchnorm updates
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad(set_to_none=True)                # clear last step's gradients
        with torch.amp.autocast(device.type):          # run forward in mixed precision
            pred = model(xb)
            loss = loss_fn(pred, yb)
        scaler.scale(loss).backward()                  # backward on the scaled loss
        scaler.unscale_(opt)                           # unscale before clipping
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # tame exploding gradients
        scaler.step(opt)                               # optimizer step
        scaler.update()                                # adjust the scale for next time

Evaluation runs on the validation split with two switches flipped. Calling model.eval() puts dropout and batch-normalization layers into inference mode, and wrapping the pass in torch.no_grad() tells autograd to build no graph, which saves memory and time. The same two switches govern the test pass and all deployment inference. After each epoch we compute the validation loss under these switches and checkpoint the model whenever it improves, so the saved file always holds the best model seen rather than the last, which may have started to overfit. The time-aware split of Section 3 guarantees this validation loss is an honest estimate of future performance.

    model.eval()                                       # inference mode for dropout and batchnorm
    val_loss = 0.0
    with torch.no_grad():                              # no autograd graph during evaluation
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            val_loss += loss_fn(model(xb), yb).item() * len(xb)
    val_loss /= len(val_loader.dataset)
    if val_loss < best_val:                            # checkpoint only on improvement
        best_val = val_loss
        torch.save({"model": model.state_dict(),       # weights only, the portable format
                    "epoch": epoch, "val_loss": val_loss}, "best.pt")

5. Inference and Export

Deployment freezes a trained module so it runs without the training stack. The first rule, repeated from Section 4, is to call model.eval() and wrap every prediction in torch.no_grad(); forgetting either inflates memory and, in the case of eval, silently changes the answer by leaving dropout active. For portability you then export the module to a self-contained graph. torch.jit.script or torch.jit.trace produces a TorchScript artifact that runs in a Python-free C++ runtime, and torch.onnx.export produces an ONNX graph that runs in cross-framework runtimes such as ONNX Runtime. Both capture the architecture and weights together, so the deployment target needs no model source code.

model.eval()
example = torch.randn(1, 16, 8, device=device)         # one window, the export's reference input

# TorchScript: a Python-free graph for the libtorch C++ runtime
scripted = torch.jit.trace(model, example)             # trace records ops on the example
scripted.save("model_ts.pt")

# ONNX: a cross-runtime graph; mark the batch and time axes as dynamic so any length runs
torch.onnx.export(model, example, "model.onnx",
                  input_names=["x"], output_names=["yhat"],
                  dynamic_axes={"x": {0: "batch", 1: "time"}, "yhat": {0: "batch"}})

Forecasting in production is often a stream: observations arrive one step at a time and you want an updated prediction after each, without re-running the model over the whole history. A recurrent model supports this naturally, because its hidden state is a sufficient summary of the past (the filtering view of Chapter 7). Stateful streaming inference, the deployment pattern detailed in Chapter 34.1, carries the hidden state between calls and feeds only the newest step, turning an $O(T)$ recompute into an $O(1)$ update per observation. The loop below shows the pattern: the LSTM consumes one step, returns a forecast, and hands its state back for the next call.

@torch.no_grad()
def stream_step(lstm, head, x_t, state):
    # x_t: (1, 1, F) one new observation; state: (h, c) carried from the previous call
    out, state = lstm(x_t, state)                      # advance the recurrence by one step
    yhat = head(out[:, -1, :])                         # forecast from the updated hidden state
    return yhat, state                                 # return state so the caller carries it forward

state = None                                           # the recurrence initializes its own state on the first call
for x_t in incoming_stream:                            # each x_t is one freshly arrived step
    yhat, state = stream_step(lstm, head, x_t, state)  # O(1) update, no recompute over history

6. JAX Notes

JAX appears throughout the structured state-space and foundation-model literature this book draws on, so it is worth a short orientation even though our code is PyTorch. JAX is built on three function transformations that compose freely. jax.jit compiles a Python function to fused XLA kernels for speed; jax.grad returns a new function that computes the gradient of a scalar-valued function; and jax.vmap vectorizes a function written for one example so it runs over a whole batch with no explicit batch loop. The defining difference from PyTorch is that JAX is functional: arrays are immutable and a model's parameters are an explicit argument passed in, rather than state hidden inside a module. The snippet below differentiates and batches a loss in the JAX style.

import jax, jax.numpy as jnp

def loss(params, x, y):                                # params passed explicitly, not stored on an object
    pred = x @ params["W"] + params["b"]
    return jnp.mean((pred - y) ** 2)

grad_fn = jax.jit(jax.grad(loss))                      # grad builds the gradient fn; jit compiles it
params = {"W": jnp.zeros((8, 4)), "b": jnp.zeros(4)}
g = grad_fn(params, jnp.ones((32, 8)), jnp.ones((32, 4)))   # gradients w.r.t. every leaf in params

batched = jax.vmap(loss, in_axes=(None, 0, 0))         # map over axis 0 of x and y, share params

JAX helps most when a temporal model exposes structure that XLA can exploit. The associative scan behind structured state-space models such as S4 and Mamba (Chapter 13) compiles to a parallel prefix sum under jit, and vmap makes per-sequence or ensemble computation trivial to express, which is why much foundation-model research ships in JAX. In practice you do not write raw JAX for a full model: Flax (or the newer Equinox) supplies the neural-network module abstraction and parameter management that nn.Module gives in PyTorch, while Optax supplies composable optimizers and gradient transforms (Adam, weight decay, gradient clipping) that you chain together and apply to the explicit parameter tree. The mental translation is direct: a PyTorch nn.Module plus torch.optim becomes a Flax module plus an Optax optimizer, with the parameters threaded through your functions instead of held inside the model.